agedi.diffusion.samplers.ffpc ============================= .. py:module:: agedi.diffusion.samplers.ffpc .. autoapi-nested-parse:: Force-field augmented predictor-corrector sampler. Classes ------- .. autoapisummary:: agedi.diffusion.samplers.ffpc.ForcefieldCorrectorSampler Module Contents --------------- .. 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).