agedi.diffusion¶

Submodules¶

Classes¶

`Agedi`	Full diffusion model: training + sampling.
`ForcefieldGuidanceConfig`	Configuration for force-field guided sampling.
`Diffusion`	Pure-Python sampling core for diffusion models.
`EulerMaruyamaSampler`	Standard Euler-Maruyama reverse-SDE sampler.
`ForcefieldCorrectorSampler`	EM predictor with force-field augmented Langevin corrector.
`HeunODESampler`	Second-order deterministic ODE sampler (Heun's method).
`HeunSampler`	Second-order stochastic sampler (Karras et al. 2022).
`PredictorCorrectorSampler`	Euler-Maruyama predictor with Langevin corrector steps.
`ProbabilityFlowODESampler`	Deterministic probability-flow ODE sampler (DDIM / Anderson 1982).
`Sampler`	Abstract base class for reverse-diffusion samplers.

Package Contents¶

class agedi.diffusion.Agedi(score_model: agedi.models.ScoreModel, noisers: List[agedi.diffusion.noisers.Noiser], regressor_model: torch.nn.Module | None = None, regressor_heads: List | None = None, regressor_loss_weight: float = 1.0, optim_config: Dict | None = None, scheduler_config: Dict | None = None, eps: float = 1e-05, fully_connected: bool = False)¶

Bases: lightning.LightningModule, agedi.diffusion.diffusion.Diffusion

Full diffusion model: training + sampling.

Combines the Diffusion sampling pipeline with LightningModule training hooks.

Parameters:

score_model (ScoreModel) – The score model.
noisers (List[Noiser]) – A list of noisers.
regressor_model (torch.nn.Module, optional) – An optional regressor model used for force-field guidance during sampling. When present, its loss is added to the diffusion loss during training.
regressor_heads (List, optional) – When provided, a RegressorModel is built internally using these heads while sharing the translator and representation from score_model. Use this parameter (instead of regressor_model) when the backbone should be shared.
regressor_loss_weight (float, optional) – Weight applied to the regressor loss. Defaults to 1.0.
optim_config (dict, optional) – Keyword arguments forwarded to torch.optim.AdamW.
scheduler_config (dict, optional) – Keyword arguments forwarded to torch.optim.lr_scheduler.ReduceLROnPlateau.
eps (float, optional) – Minimum diffusion time value.

regressor_loss_weight = 1.0¶

optim_config = None¶

scheduler_config = None¶

_regressor_training = False¶

fully_connected = False¶

on_fit_start() → None¶: Write hparams.yaml to the trainer log directory at training start.

get_hparams() → Dict¶

Return hyperparameters sufficient to reconstruct this diffusion model.

Returns:: Hyperparameter dictionary with _target_, score_model, noisers, optim_config, scheduler_config, eps, and optionally regressor_heads or regressor_model.
Return type:: dict

setup(stage: str = None) → None¶: Set up the model (put score model in training mode).

forward(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Forward pass through the score model.

Parameters:: batch (AtomsGraph) – A batch of AtomsGraph data.
Returns:: The output of the score model forward pass.
Return type:: AtomsGraph

loss(batch: agedi.data.AtomsGraph, batch_idx: torch.Tensor) → Dict¶

Compute the combined diffusion + regressor loss.

Always computes the diffusion (denoising) loss on a noised copy of the batch. When a regressor model is present and the batch contains force labels, the regressor loss is added with weight regressor_loss_weight.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
batch_idx (torch.Tensor) – The index of the batch.

Returns:

A dictionary of losses.

Return type:

dict

diffusion_loss(batch: agedi.data.AtomsGraph, batch_idx: torch.Tensor) → Dict¶

Compute the diffusion (denoising score-matching) loss.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
batch_idx (torch.Tensor) – The index of the batch.

Returns:

A dictionary of losses.

Return type:

dict

regressor_loss(batch: agedi.data.AtomsGraph, batch_idx: torch.Tensor) → Dict¶

Compute the regressor loss on the un-noised batch.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
batch_idx (torch.Tensor) – The index of the batch.

Returns:

A dictionary of losses.

Return type:

dict

Raises:

ValueError – If no regressor model is attached.

training_step(batch, batch_idx: torch.Tensor) → torch.Tensor¶

Perform a training step.

Computes the combined diffusion + regressor loss (see loss()).

When the Dataset was set up with a dedicated regressor dataset (via add_regressor_data()), batch is a dict with two keys:

"main" – a regular training batch used for both the diffusion and regressor loss.
"regressor" – a regressor-only batch whose structures are only forwarded through the regressor loss (not the diffusion loss).

When no regressor dataset is present batch is a plain AtomsGraph batch and the behaviour is identical to the pre-existing implementation.

Parameters:

batch (AtomsGraph or dict) – A batch of AtomsGraph data, or a dict with "main" and "regressor" keys when a dedicated regressor dataset is used.
batch_idx (torch.Tensor) – The index of the batch.

Returns:

The combined loss.

Return type:

torch.Tensor

validation_step(batch: agedi.data.AtomsGraph, batch_idx: torch.Tensor) → torch.Tensor¶

Perform a validation step.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
batch_idx (torch.Tensor) – The index of the batch.

Returns:

The combined loss.

Return type:

torch.Tensor

configure_optimizers() → Dict¶

Configure optimizers and learning-rate schedulers.

When a regressor model is present a single optimizer is built over the deduplicated union of score_model and regressor_model parameters (shared parameters appear only once).

Returns:: A dictionary with "optimizer", "lr_scheduler", and "monitor" keys.
Return type:: dict

_scheduler_monitor() → str¶: Return the metric used by ReduceLROnPlateau.

property regressor_training: bool¶: Whether the regressor model is in training mode.

class agedi.diffusion.ForcefieldGuidanceConfig¶

Configuration for force-field guided sampling.

Parameters:

guidance (float) – Scale of the force-field guidance applied at each reverse step. Set to 0.0 (the default) to disable guidance entirely.
zeta (float) – Exponent for the time-dependent weight factor (1 - t)**zeta. Higher values concentrate guidance near the end of the trajectory.
force_threshold (float) – Convergence criterion for the optional post-diffusion relaxation: the maximum per-atom force magnitude (eV/Å) below which relaxation stops.
max_extra_steps (int) – Maximum number of additional relaxation steps performed after the main diffusion trajectory when guidance > 0.

guidance: float = 0.0¶

zeta: float = 3.0¶

force_threshold: float = 0.05¶

max_extra_steps: int = 0¶

class agedi.diffusion.Diffusion(score_model: ScoreModel, noisers: List[agedi.diffusion.noisers.Noiser], regressor_model: torch.nn.Module | None = None, eps: float = 1e-05)¶

Pure-Python sampling core for diffusion models.

Holds the score model, noisers, and an optional regressor and provides the full forward / reverse / sampling pipeline. This class does not inherit from torch.nn.Module or lightning.LightningModule and therefore has no training hooks.

When used through Agedi (which inherits from both this class and lightning.LightningModule), the Lightning infrastructure manages device placement and module registration. When used standalone, device information is derived from the score model’s parameters via the device property.

Parameters:

score_model (ScoreModel) – The score model.
noisers (List[Noiser]) – A list of noisers.
regressor_model (torch.nn.Module, optional) – An optional regressor model used for force-field guidance during sampling.
eps (float, optional) – Minimum value for the diffusion time step (used in sample_time()).

score_model¶

noisers¶

regressor_model = None¶

eps = 1e-05¶

lbfgs_step_sizer: agedi.diffusion.guidance.BatchedLBFGSStepSizer | None = None¶

zeta: float = 3.0¶

noiser_keys¶

score_keys¶

_compiled_reverse_step = None¶

property device: torch.device¶

Infer the computation device from the score model’s parameters.

When used through Agedi (which also inherits lightning.LightningModule), Lightning’s own device property takes precedence.

sample_time(batch: agedi.data.AtomsGraph) → None¶

Sample a random diffusion time for each graph in batch.

Draws times uniformly from [eps, 1] and assigns them to batch.time at atom resolution.

Parameters:: batch (AtomsGraph) – A batch of AtomsGraph data; modified in-place.

forward_step(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Forward diffusion step (corruption).

Applies each noiser in order to corrupt the batch.

Parameters:: batch (AtomsGraph) – A batch of AtomsGraph data.
Returns:: The corrupted batch.
Return type:: AtomsGraph

reverse_step(batch: agedi.data.AtomsGraph, delta_t: float, force_field_guidance: float, last: bool = False, timings: SamplingTimings | None = None) → agedi.data.AtomsGraph¶

Reverse diffusion step (denoising).

Evaluates the score model and applies one reverse-SDE step through all noisers. Optionally applies force-field guidance afterwards.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
delta_t (float) – The time step.
force_field_guidance (float) – Scale of the force-field guidance (0.0 disables it).
last (bool, optional) – Whether this is the final denoising step.
timings (SamplingTimings, optional) – If provided, timing measurements are accumulated here.

Returns:

The denoised batch.

Return type:

AtomsGraph

_resolve_sampler(sampler: str | Sampler | None, corrector_steps: int = 0, corrector_step_size: float = 0.001, sampler_kwargs: Dict | None = None) → Sampler¶

Resolve a sampler string/instance or build one from legacy params.

Parameters:

sampler (str, Sampler, or None) – None → EulerMaruyamaSampler (or PredictorCorrectorSampler when corrector_steps > 0). A string looks up the sampler in the registry. A Sampler instance is returned as-is.
corrector_steps (int, optional) – Used when sampler is None and corrector_steps > 0 to build a PredictorCorrectorSampler.
corrector_step_size (float, optional) – Step size forwarded to the predictor-corrector sampler.
sampler_kwargs (dict, optional) – Extra keyword arguments forwarded to the sampler constructor when sampler is a string alias. Keys override the defaults supplied by corrector_steps / corrector_step_size.

Returns:

A ready-to-use sampler instance.

Return type:

Sampler

corrector_step(batch: agedi.data.AtomsGraph, corrector_dt: float) → agedi.data.AtomsGraph¶

Langevin corrector step at constant time.

Evaluates the score model and applies one Langevin corrector step through all noisers (in reverse order).

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
corrector_dt (float) – Step size for the Langevin corrector.

Returns:

The corrected batch.

Return type:

AtomsGraph

force_field_guidance_step(batch: agedi.data.AtomsGraph, scale: float, max_step_size: float = 0.1) → agedi.data.AtomsGraph¶

Apply one force-field guidance step.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
scale (float) – Base scale of the force field guidance.
max_step_size (float, optional) – Maximum allowed step size magnitude.

Returns:

Updated batch.

Return type:

AtomsGraph

post_diffusion_relaxation_step(batch: agedi.data.AtomsGraph, scale: float = 0.1) → agedi.data.AtomsGraph¶

Perform a pure force-based relaxation step.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
scale (float, optional) – Step size scaling factor.

Returns:

Updated batch.

Return type:

AtomsGraph

_initialize_graph(cutoff: float, fully_connected: bool = False, **kwargs) → agedi.data.AtomsGraph¶

Initialise a single graph from noiser priors.

Parameters:

cutoff (float) – Cutoff radius for the neighbour list.
fully_connected (bool, optional) – When True the graph is rebuilt as a fully connected graph at every reverse step instead of using a finite cutoff. Recommended for gas-phase molecules and clusters. Defaults to False.
**kwargs – Additional keyword arguments passed to the graph (e.g. cell, template, pbc).

Returns:

The initialised graph.

Return type:

AtomsGraph

static _sync_for_timing(device: torch.device | None) → None¶

_time_sampling_call(device: torch.device | None, timings: SamplingTimings, key: str, fn, *args, **kwargs)¶

static _format_timing_line(label: str, value: float, count: int | None = None) → str¶

_print_sampling_timings(timings: SamplingTimings) → None¶

property compiled_reverse_step¶

Lazily compile reverse_step() with torch.compile.

The compiled kernel is cached as self._compiled_reverse_step so that compilation happens at most once per model instance. Using a per-instance cache (rather than a class-level @torch.compile decorator) means that two Diffusion objects with different architectures will each compile their own kernel and never interfere.

Note

timings must not be passed to the compiled function — time.perf_counter is not traceable by Dynamo. Time the compiled call from outside in _sample_batch() using the is_compiled flag.

_sample_batch(batch: torch_geometric.data.Batch, steps: int, eps: float, force_field_guidance: float, save_trajectory: bool, progress_bar: bool, force_threshold: float, max_extra_steps: int, corrector_steps: int = 0, corrector_step_size: float = 0.001, timings: SamplingTimings | None = None, reverse_step_fn=None, is_compiled: bool = False, sampler=None, sampler_kwargs=None) → List[agedi.data.AtomsGraph]¶

Run the reverse-diffusion loop for a pre-built batch.

Parameters:

batch (Batch) – A batch of AtomsGraph data at t=1.
steps (int) – Number of reverse-diffusion steps.
eps (float) – Minimum time value (end of trajectory).
force_field_guidance (float) – Scale of the force-field guidance (0.0 disables it).
save_trajectory (bool) – Whether to collect and return all intermediate states.
progress_bar (bool) – Whether to display a tqdm progress bar.
force_threshold (float) – Maximum per-atom force for terminating post-diffusion relaxation.
max_extra_steps (int) – Maximum extra relaxation steps after the main trajectory.
corrector_steps (int, optional) – Number of Langevin corrector passes after each predictor step. 0 (default) disables the corrector (standard DDPM/EM sampling).
corrector_step_size (float, optional) – Step size used for each Langevin corrector step. Defaults to 1e-3.
timings (SamplingTimings, optional) – If provided, timing measurements are accumulated here.
reverse_step_fn (callable, optional) – The reverse step function to use. Defaults to self.reverse_step. Pass a torch.compile-wrapped version to enable compiled sampling.
is_compiled (bool, optional) – Whether reverse_step_fn is a compiled function.
sampler (str, Sampler, or None, optional) – Sampler instance or string alias controlling the reverse-diffusion algorithm. When provided (and is_compiled is False), the sampler’s step() is called instead of reverse_step_fn. None (default) falls back to an EulerMaruyamaSampler (or a PredictorCorrectorSampler when corrector_steps > 0).
sampler_kwargs (dict, optional) – Extra constructor arguments forwarded to the sampler when sampler is a string alias. Keys override the defaults supplied by corrector_steps / corrector_step_size.

Returns:

Final structures, or (when save_trajectory is True) a list of trajectories (one per graph).

Return type:

List[AtomsGraph]

_sample(N: int, steps: int, cutoff: float, eps: float, force_field_guidance: float, force_threshold: float, max_extra_steps: int, progress_bar: bool, save_trajectory: bool, corrector_steps: int = 0, corrector_step_size: float = 0.001, print_timings: bool = False, compile: bool = False, sampler=None, sampler_kwargs=None, **kwargs) → List[agedi.data.AtomsGraph]¶

Build N graphs from priors and run the sampling loop.

Parameters:

N (int) – Number of structures to generate.
steps (int) – Number of reverse-diffusion steps.
cutoff (float) – Cutoff radius for the neighbour list.
eps (float) – Minimum time value (end of trajectory).
force_field_guidance (float) – Scale of the force-field guidance.
force_threshold (float) – Maximum per-atom force for post-diffusion relaxation.
max_extra_steps (int) – Maximum extra relaxation steps.
progress_bar (bool) – Show tqdm progress bar.
save_trajectory (bool) – Collect all intermediate states.
corrector_steps (int, optional) – Langevin corrector passes per predictor step.
corrector_step_size (float, optional) – Step size for each corrector pass.
print_timings (bool, optional) – Print a timing breakdown after sampling completes.
compile (bool, optional) – Use torch.compile on the reverse diffusion step.
sampler_kwargs (dict, optional) – Extra keyword arguments forwarded to the sampler constructor when sampler is a string alias.
**kwargs – Keyword arguments forwarded to _initialize_graph().

Returns:

Sampled structures (or trajectories when save_trajectory is True).

Return type:

List[AtomsGraph]

sample(N: int, template=None, batch_size: int | None = 64, steps: int | None = 500, cutoff: float | None = 6.0, eps: float | None = 0.001, n_atoms: int | None = None, atomic_numbers: List[int] | None = None, formula: str | None = None, positions: numpy.ndarray | None = None, cell: numpy.ndarray | None = None, pbc: numpy.ndarray | None = None, confinement: Tuple[float, float] | None = None, compile: bool = False, ff_guidance: agedi.diffusion.guidance.ForcefieldGuidanceConfig | None = None, property: Dict | None = None, progress_bar: bool | None = False, save_trajectory: bool | None = False, print_timings: bool | None = False, corrector_steps: int = 0, corrector_step_size: float = 0.001, sampler=None, sampler_kwargs=None) → List[agedi.data.AtomsGraph]¶

Sample structures from the diffusion model.

The minimum required arguments depend on the configured noisers and whether a template is provided:

n_atoms – always required unless derivable from atomic_numbers or formula.
atomic_numbers – required unless a types-noiser is configured (key "x"), or derivable from formula.
positions – required when no positions-noiser is configured (type-only diffusion).
cell – required for periodic systems when no template is given. Not required when pbc=[False, False, False].
pbc – optional; defaults to [True, True, True]. Pass [False, False, False] for non-periodic systems.

Parameters:

N (int) – Number of structures to generate.
template (AtomsGraph or ase.Atoms, optional) – Template structure. cell and pbc are taken from the template when not explicitly provided.
batch_size (int, optional) – Internal batch size for splitting large N.
steps (int, optional) – Number of reverse-diffusion steps.
cutoff (float, optional) – Cutoff radius for the neighbour list.
eps (float, optional) – Minimum time value at the end of the trajectory.
n_atoms (int, optional) – Number of atoms per structure.
atomic_numbers (List[int], optional) – Atomic numbers of the atoms to generate.
formula (str, optional) – Chemical formula (e.g. "H2O").
positions (np.ndarray, optional) – Fixed atom positions (shape (n_atoms, 3)).
cell (np.ndarray, optional) – Unit-cell matrix (3x3).
pbc (np.ndarray, optional) – Periodic boundary conditions.
confinement (Tuple[float, float], optional) – Z-directional confinement (z_min, z_max).
compile (bool, optional) – When True, use torch.compile on the reverse diffusion step for improved throughput on CUDA hardware.
ff_guidance (ForcefieldGuidanceConfig, optional) – Force-field guidance configuration.
property (dict, optional) – Conditioning properties (key -> scalar tensor).
progress_bar (bool, optional) – Show a tqdm progress bar.
save_trajectory (bool, optional) – Return full trajectories instead of final structures.
print_timings (bool, optional) – Print a timing breakdown after sampling completes.
corrector_steps (int, optional) – Number of Langevin corrector passes after each predictor step. 0 (default) gives standard (predictor-only) sampling. Ignored when sampler is provided explicitly.
corrector_step_size (float, optional) – Step size for each corrector pass. Defaults to 1e-3.
sampler (str, Sampler, or None, optional) –
Reverse-diffusion algorithm. Pass a string alias or a Sampler instance.

Available string aliases:
- "em" — Euler-Maruyama (default)
- "pc" — Predictor-corrector (use with corrector_steps)
- "heun" — Stochastic Heun, 2nd-order (2 score calls/step)
- "ddim" — Deterministic probability-flow ODE (DDIM)
- "heun_ode" — Deterministic 2nd-order ODE (Heun on PF-ODE)
- "ffpc" — Force-field corrector (use with sampler_kwargs)
When None (default), uses "em" (or "pc" when corrector_steps > 0).

sampler_kwargs (dict, optional) –

Sampler-specific constructor arguments, forwarded when sampler is a string alias. Keys override the top-level corrector_steps and corrector_step_size defaults. Examples:

# PC with 3 corrector steps and a custom step size
model.sample(N, sampler="pc",
             sampler_kwargs={"corrector_steps": 3,
                            "corrector_step_size": 5e-4})

# FFPC with 5 force-field corrector steps
model.sample(N, sampler="ffpc",
             sampler_kwargs={"corrector_steps": 5,
                            "corrector_scale": 0.005})

Returns:

Sampled structures, or trajectories when save_trajectory is True.

Return type:

List[AtomsGraph]

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

Bases: 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 denoise() method. The update formula used (EM or DDPM posterior mean) is controlled by the sampler attribute of each PositionsNoiser.

This is the default sampler and exactly reproduces the behaviour of reverse_step() (minus guidance and timings).

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶

Euler-Maruyama reverse step.

Evaluate score model.
Apply each noiser’s denoising update in reverse order.
Wrap positions and rebuild the neighbour list.

class agedi.diffusion.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.

EM predictor with neural score.
Advance batch.time to t - dt.
For each corrector iteration: evaluate neural score, blend with force-field gradient, apply Langevin step.
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).

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

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

First score at (x_t, t): s_1 = s_θ(x_t, t)
ODE predictor: x̃ = x_t + dt·(0.5·g(t)²·s_1 - f(x_t, t))
Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)
Averaged score: s_avg = 0.5·(s_1 + s_2)
ODE corrector: x_{t-dt} = x_t + dt·(0.5·g(t)²·s_avg - f(x_t, t))

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

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

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.

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

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

First score at (x_t, t): s_1 = s_θ(x_t, t)
Deterministic EM predictor (no noise) for SDE noisers: x̃ = x_t + dt·(g(t)²·s_1 - f(x_t, t))
Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)
Average scores: s_avg = 0.5·(s_1 + s_2)
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 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.

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

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.

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶: Second-order stochastic Heun step.

class agedi.diffusion.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: 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).

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.
corrector_steps (int, optional) – Number of Langevin corrector passes per predictor step. Default: 1.
corrector_step_size (float, optional) – Step size for each Langevin corrector step. Default: 1e-3.

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 sample() uses the same convention via this class.

corrector_steps = 1¶

corrector_step_size = 0.001¶

_corrector_dt: torch.Tensor | None = None¶

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶

Predictor-corrector step.

EM predictor: evaluate score, apply noiser.denoise(), wrap & rebuild.
Advance batch.time to t - dt.
For each corrector iteration: evaluate score, apply Langevin step via noiser.langevin_step(), wrap & rebuild.

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

Bases: agedi.diffusion.samplers.base.Sampler

Deterministic probability-flow ODE sampler (DDIM / Anderson 1982).

Integrates the reverse-time ODE:

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

\[\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 EulerMaruyamaSampler, this sampler is fully deterministic: repeated calls with identical inputs produce identical outputs.

Parameters:

score_fn (callable) – Score-model forward function.
noisers (list of Noiser) – Noisers in forward order.

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶

Probability-flow ODE step.

Evaluate score model.
For each SDE-based noiser: apply the ODE update (half diffusion, no noise).
For other noisers: apply a deterministic EM step.
Wrap positions and rebuild the neighbour list.

The last parameter is accepted for API compatibility but has no effect — this sampler is always deterministic.

static _ode_step(batch: agedi.data.AtomsGraph, noiser: agedi.diffusion.noisers.Noiser, dt: torch.Tensor) → agedi.data.AtomsGraph¶

Apply the probability-flow ODE update for one SDE-based noiser.

Computes:

\[\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)\]

Parameters:

batch (AtomsGraph) – Current batch state.
noiser (Noiser) – An SDE-based noiser with a .sde attribute.
dt (torch.Tensor) – Step size.

Returns:

Batch with the updated attribute.

Return type:

AtomsGraph

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

Bases: 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 _sample_batch() sets batch.time before each call to step() and handles trajectory saving, force-field guidance, and post-diffusion relaxation. The sampler is responsible only for the diffusion update itself.

Parameters:

score_fn (callable) – A function score_fn(batch) -> batch that runs the score model and stores per-key scores in batch.{key}_score.
noisers (list of Noiser) – The noisers attached to the diffusion model, in forward order. The sampler iterates them in reverse when denoising.
Registry
--------

:param Use 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" — EulerMaruyamaSampler: :param * "pc" — PredictorCorrectorSampler: :param * "heun" — HeunSampler: :param * "ddim" — ProbabilityFlowODESampler: :param * "heun_ode" — HeunODESampler: :param * "ffpc" — ForcefieldCorrectorSampler:

_registry: ClassVar[Dict[str, Callable[Ellipsis, Sampler]]]¶

uses_force_field: ClassVar[bool] = False¶

score_fn¶

noisers¶

classmethod register(name: str, factory: Callable[Ellipsis, Sampler]) → None¶

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

Parameters:

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

Examples

from agedi.diffusion.samplers import Sampler

class MySampler(Sampler):
    ...

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

static _check_finite(batch: agedi.data.AtomsGraph, stage: str) → None¶

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.

Parameters:

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

abstractmethod step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph¶

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.

Parameters:

batch (AtomsGraph) – Current state of the batch. batch.time is the current time.
dt (torch.Tensor) – Positive step size: dt = t_i - t_{i+1}.
last (bool) – 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).

Returns:

Updated batch with a valid neighbour list.

Return type:

AtomsGraph