agedi.diffusion.samplers.ffpc

Force-field augmented predictor-corrector sampler.

Classes

ForcefieldCorrectorSampler

EM predictor with force-field augmented Langevin corrector.

Module Contents

class agedi.diffusion.samplers.ffpc.ForcefieldCorrectorSampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser], regressor_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph] | None = None, corrector_steps: int = 1, corrector_step_size: float = 0.001, mixing_zeta: float = 1.0, temperature: float | None = None, terminal_steps: int = 0, terminal_step_size: float | None = None, terminal_dynamics: Literal['overdamped', 'langevin_md'] = 'overdamped', terminal_friction: float | None = None)

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

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

\[\tilde{s}(x, t) = (1 - f(t))\, s_\theta(x) + f(t)\, F(x)\]

where \(f(t) = (1 - t)^\zeta\) increases as \(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):

\[x_{n+1} = x_n + \varepsilon\, F(x_n) + \sqrt{2\varepsilon T}\, z\]

where \(\varepsilon = \texttt{terminal\_step\_size}\) is a reduced step \(\varepsilon = \Delta t / T\). This parameterization gives T-independent stability (condition: \(\varepsilon < 2/k_{\max}\)) while preserving the correct Boltzmann stationary distribution \(\propto \exp(-U / T)\) through the \(\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:

\[p_i \sim \mathcal{N}\!\left(0,\, m_i T\right)\]

followed by BAOAB integration:

\[\begin{split}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\end{split}\]

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.

Parameters:

  • score_fn (callable) – Neural score model: score_fn(batch) -> batch with pos_score set.

  • noisers (list[Noiser]) – Noisers in forward order; iterated in reverse for denoising.

  • regressor_fn (callable or None) – Force-field model: regressor_fn(batch) -> batch with forces_prediction set. Passed by _resolve_sampler().

  • corrector_steps (int) – Number of Langevin corrector steps per predictor step. Default: 1.

  • corrector_step_size (float) – Step size for each Langevin corrector step. Subject to the Langevin stability bound corrector_step_size < 2·var(t_{i-1}). Default: 1e-3.

  • mixing_zeta (float) – 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.

  • temperature (float) –

    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 \(\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 \(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 UserWarning reminding you to set an explicit value.

  • terminal_steps (int) – Number of refinement steps applied after the last diffusion step. 0 (default) disables them.

  • terminal_step_size (float or None) –

    Step size for each terminal step.

    • overdamped: the reduced step \(\varepsilon = \Delta t / T\). Stability requires \(\varepsilon < 2 / k_{\max}\) where \(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.

  • terminal_dynamics ("overdamped" or "langevin_md") – 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.

  • terminal_friction (float or None) – Friction coefficient \(\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 \(\gamma = 0.1 / \Delta t\), giving a dimensionless reduced friction \(\gamma \Delta t = 0.1\) (moderately damped).

  • alias (String)

  • ------------

:param "ffpc" — registered in agedi.diffusion.samplers.:

Notes

This sampler calls regressor_fn directly, not force_field_guidance_step. The LBFGS step sizer is therefore not required; uses_force_field is False.

uses_force_field: ClassVar[bool] = False
regressor_fn = None
corrector_steps = 1
corrector_step_size = 0.001
mixing_zeta = 1.0
temperature = None
terminal_steps = 0
terminal_step_size = None
terminal_dynamics = 'overdamped'
terminal_friction = None
_pos_noiser
_corrector_dt: torch.Tensor | None = None
_pending_frames: List = []
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 _pending_frames for trajectory capture.

_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.

_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.

_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 \(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).