agedi.diffusion.samplers ======================== .. py:module:: agedi.diffusion.samplers .. autoapi-nested-parse:: Pluggable reverse-diffusion samplers for AGeDi. Each sampler encapsulates one complete reverse-diffusion outer step, from time ``t`` to ``t - dt``. Samplers may call the score model (and optionally the force-field model) any number of times within that outer step. Available samplers ------------------ :class:`EulerMaruyamaSampler` Standard Euler-Maruyama reverse-SDE step (the default). One score call per step. String alias: ``"em"``. :class:`PredictorCorrectorSampler` EM predictor followed by ``N`` Langevin corrector steps at the new noise level. String alias: ``"pc"``. :class:`HeunSampler` Second-order stochastic sampler (Karras et al. 2022). Two score calls per step; the averaged score is used in the final stochastic update. String alias: ``"heun"``. :class:`ProbabilityFlowODESampler` Deterministic probability-flow ODE sampler (DDIM / Anderson 1982). One score call per step; applies half the diffusion coefficient with no noise. String alias: ``"ddim"``. :class:`HeunODESampler` Second-order deterministic ODE sampler (Heun's method on the PF-ODE). Two score calls per step; fully deterministic. String alias: ``"heun_ode"``. Usage ----- Pass a sampler by string alias or instance to :meth:`~agedi.diffusion.Diffusion.sample`:: structures = model.sample(10, steps=200, sampler="heun") structures = model.sample(10, steps=500, sampler=PredictorCorrectorSampler(model.score_model, model.noisers, corrector_steps=3)) Register a custom sampler:: from agedi.diffusion.samplers import Sampler class MySampler(Sampler): def step(self, batch, dt, last): ... Sampler.register("my_sampler", lambda score_fn, noisers, **kw: MySampler(score_fn, noisers)) Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/agedi/diffusion/samplers/base/index /autoapi/agedi/diffusion/samplers/em/index /autoapi/agedi/diffusion/samplers/ffpc/index /autoapi/agedi/diffusion/samplers/heun/index /autoapi/agedi/diffusion/samplers/ode/index /autoapi/agedi/diffusion/samplers/pc/index Classes ------- .. autoapisummary:: agedi.diffusion.samplers.Sampler agedi.diffusion.samplers.EulerMaruyamaSampler agedi.diffusion.samplers.ForcefieldCorrectorSampler agedi.diffusion.samplers.HeunSampler agedi.diffusion.samplers.HeunODESampler agedi.diffusion.samplers.ProbabilityFlowODESampler agedi.diffusion.samplers.PredictorCorrectorSampler Package 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 .. 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. .. py:class:: ForcefieldCorrectorSampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser], regressor_fn: Optional[Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph]] = None, corrector_steps: int = 1, corrector_step_size: float = 0.001, mixing_zeta: float = 1.0, temperature: Optional[float] = None, terminal_steps: int = 0, terminal_step_size: Optional[float] = None, terminal_dynamics: Literal['overdamped', 'langevin_md'] = 'overdamped', terminal_friction: Optional[float] = None) Bases: :py:obj:`agedi.diffusion.samplers.base.Sampler` EM predictor with force-field augmented Langevin corrector. Follows the standard predictor-corrector scheme from Song et al. (2021) but uses a force-field augmented score in the corrector steps: **Predictor** — standard EM step using the neural score only: .. math:: x_{t-\Delta t} = x_t + \Delta t\,(g(t)^2\, s_\theta(x_t) - f(x_t,\,t)) + \sqrt{\Delta t}\, g(t)\, z **Corrector** — ``corrector_steps`` Langevin steps at the new noise level ``t - \Delta t``, using a score that blends the neural model with the force-field gradient: .. math:: \tilde{s}(x, t) = (1 - f(t))\, s_\theta(x) + f(t)\, F(x) where :math:`f(t) = (1 - t)^\zeta` increases as :math:`t \to 0`, so the force field has no influence at the start of sampling and full influence at the end. Temperature is **not** applied here; the relative weight of the force field is controlled entirely by ``mixing_zeta`` and ``corrector_step_size``. **Terminal phase** (optional) — applies ``terminal_steps`` additional refinement steps after the last diffusion step using only the force field. Two dynamics modes are supported via ``terminal_dynamics``: *overdamped* (default) — overdamped Langevin (no momenta): .. math:: x_{n+1} = x_n + \varepsilon\, F(x_n) + \sqrt{2\varepsilon T}\, z where :math:`\varepsilon = \texttt{terminal\_step\_size}` is a **reduced step** :math:`\varepsilon = \Delta t / T`. This parameterization gives T-independent stability (condition: :math:`\varepsilon < 2/k_{\max}`) while preserving the correct Boltzmann stationary distribution :math:`\propto \exp(-U / T)` through the :math:`\sqrt{T}` factor in the noise amplitude. *langevin_md* — standard Langevin MD (BAOAB integrator) with atomic masses from ASE. Momenta are initialised from the Maxwell-Boltzmann distribution at temperature *T*: .. math:: p_i \sim \mathcal{N}\!\left(0,\, m_i T\right) followed by BAOAB integration: .. math:: p &\leftarrow p + \tfrac{\Delta t}{2}\, F \\ x &\leftarrow x + \tfrac{\Delta t}{2}\, p / m \\ p &\leftarrow e^{-\gamma \Delta t}\, p + \sqrt{T\, m\,(1 - e^{-2\gamma \Delta t})}\, z \\ x &\leftarrow x + \tfrac{\Delta t}{2}\, p / m \\ F &\leftarrow F(x) \\ p &\leftarrow p + \tfrac{\Delta t}{2}\, F When no regressor is attached (``regressor_fn=None``), the corrector falls back to a pure Langevin step with the neural score, and terminal steps are skipped. :param score_fn: Neural score model: ``score_fn(batch) -> batch`` with ``pos_score`` set. :type score_fn: callable :param noisers: Noisers in forward order; iterated in reverse for denoising. :type noisers: list[Noiser] :param regressor_fn: Force-field model: ``regressor_fn(batch) -> batch`` with ``forces_prediction`` set. Passed by :meth:`~agedi.diffusion.Diffusion._resolve_sampler`. :type regressor_fn: callable or None :param corrector_steps: Number of Langevin corrector steps per predictor step. Default: ``1``. :type corrector_steps: int :param corrector_step_size: Step size for each Langevin corrector step. Subject to the Langevin stability bound ``corrector_step_size < 2·var(t_{i-1})``. Default: ``1e-3``. :type corrector_step_size: float :param mixing_zeta: Exponent for the mixing schedule ``f(t) = (1 - t)**mixing_zeta``. ``1.0`` (default) gives linear mixing; higher values concentrate force-field influence near the end of the trajectory. :type mixing_zeta: float :param temperature: Temperature *T* used **only** in the terminal phase (not in the corrector steps, where dividing forces by T causes blowup at low T). * *overdamped*: appears as :math:`\sqrt{T}` in the noise amplitude (see the update rule above) to preserve the correct Boltzmann distribution. The step-size stability bound is T-independent. * *langevin_md*: the thermal energy :math:`k_B T` in the same units as the model forces (typically eV). Sets the noise amplitude in the Ornstein-Uhlenbeck step and the Maxwell-Boltzmann velocity initialisation. Default: ``None`` — falls back to ``1.0`` with a :class:`UserWarning` reminding you to set an explicit value. :type temperature: float :param terminal_steps: Number of refinement steps applied after the last diffusion step. ``0`` (default) disables them. :type terminal_steps: int :param terminal_step_size: Step size for each terminal step. * *overdamped*: the **reduced step** :math:`\varepsilon = \Delta t / T`. Stability requires :math:`\varepsilon < 2 / k_{\max}` where :math:`k_{\max}` is the largest force-constant eigenvalue — a bound that is independent of *T*. ``None`` (default) auto-selects ``1e-3``, which is conservative for most force fields. * *langevin_md*: physical time step in units consistent with the model forces and ASE masses. When forces are in eV/Å and masses in amu the unit is femtoseconds. ``None`` (default) auto-selects ``1.0`` (1 fs), a safe starting point for typical ML potentials. :type terminal_step_size: float or None :param terminal_dynamics: Dynamics used for the terminal phase. ``"overdamped"`` (default) uses overdamped Langevin (no momenta). ``"langevin_md"`` uses standard Langevin MD (BAOAB) with real atomic masses from ASE and momenta initialised from the Maxwell-Boltzmann distribution. :type terminal_dynamics: ``"overdamped"`` or ``"langevin_md"`` :param terminal_friction: Friction coefficient :math:`\gamma` for the Langevin thermostat in *langevin_md* dynamics (ignored for *overdamped*). Units are the inverse of ``terminal_step_size``; when ``terminal_step_size`` is in fs a physically reasonable range is ``0.001``–``0.1`` fs⁻¹ (1–100 ps⁻¹). ``None`` (default) auto-selects :math:`\gamma = 0.1 / \Delta t`, giving a dimensionless reduced friction :math:`\gamma \Delta t = 0.1` (moderately damped). :type terminal_friction: float or None :param String alias: :param ------------: :param ``"ffpc"`` — registered in :mod:`agedi.diffusion.samplers`.: .. rubric:: Notes This sampler calls :attr:`regressor_fn` directly, *not* ``force_field_guidance_step``. The LBFGS step sizer is therefore not required; ``uses_force_field`` is ``False``. .. py:attribute:: uses_force_field :type: ClassVar[bool] :value: False .. py:attribute:: regressor_fn :value: None .. py:attribute:: corrector_steps :value: 1 .. py:attribute:: corrector_step_size :value: 0.001 .. py:attribute:: mixing_zeta :value: 1.0 .. py:attribute:: temperature :value: None .. py:attribute:: terminal_steps :value: 0 .. py:attribute:: terminal_step_size :value: None .. py:attribute:: terminal_dynamics :value: 'overdamped' .. py:attribute:: terminal_friction :value: None .. py:attribute:: _pos_noiser .. py:attribute:: _corrector_dt :type: Optional[torch.Tensor] :value: None .. py:attribute:: _pending_frames :type: List :value: [] .. py:method:: step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) -> agedi.data.AtomsGraph Predictor-corrector step with force-field augmented corrector. 1. EM predictor with neural score. 2. Advance ``batch.time`` to ``t - dt``. 3. For each corrector iteration: evaluate neural score, blend with force-field gradient, apply Langevin step. 4. If ``last`` and ``terminal_steps > 0``: terminal phase with the chosen dynamics. Intermediate frames are stored in :attr:`_pending_frames` for trajectory capture. .. py:method:: _run_terminal(batch: agedi.data.AtomsGraph) -> None Dispatch to the selected terminal dynamics. Prepends a bridge frame (the denoised state before terminal dynamics start) unless corrector frames already captured it as their last entry — which is the case when ``corrector_steps > 0``. .. py:method:: _terminal_overdamped(batch: agedi.data.AtomsGraph) -> None Overdamped Langevin terminal steps (no momenta). ``terminal_step_size`` is the reduced step ε = dt/T. Update: x += ε·F + sqrt(2εT)·z, giving stationary ∝ exp(-U/T). Stability condition: ε < 2/k_max, independent of T. .. py:method:: _terminal_langevin_md(batch: agedi.data.AtomsGraph) -> None Standard Langevin MD terminal steps (BAOAB, real atomic masses). Masses are looked up from ASE using ``batch.x`` (atomic numbers). Temperature *T* is the thermal energy :math:`k_B T` in model force units. Time step units must be consistent with forces and masses (e.g. fs when forces are in eV/Å and masses in amu). .. 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. .. 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. .. 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:: 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.