agedi.diffusion.guidance¶

Force-field guidance utilities for diffusion sampling.

This module provides: - ForcefieldGuidanceConfig – configuration dataclass. - LBFGSStepSizer – per-graph L-BFGS step-size adapter. - BatchedLBFGSStepSizer – batched wrapper around LBFGSStepSizer. - force_field_guidance_step() – one guidance step (module-level). - post_diffusion_relaxation_step() – post-diffusion relaxation (module-level).

Classes¶

`ForcefieldGuidanceConfig`	Configuration for force-field guided sampling.
`LBFGSStepSizer`	L-BFGS approach for determining optimal step sizes in force field guidance.
`BatchedLBFGSStepSizer`	Batched wrapper around `LBFGSStepSizer` for use with batched graphs.

Functions¶

`force_field_guidance_step`(→ agedi.data.AtomsGraph)	Apply one force-field guidance step with batched L-BFGS step-size adaptation.
`post_diffusion_relaxation_step`(→ agedi.data.AtomsGraph)	Perform a pure force-based relaxation step after diffusion is complete.

Module Contents¶

class agedi.diffusion.guidance.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.guidance.LBFGSStepSizer(memory_size: int = 10, initial_step: float = 0.1, device: str = 'cuda')¶

L-BFGS approach for determining optimal step sizes in force field guidance.

memory_size = 10¶

initial_step = 0.1¶

device = 'cuda'¶

s_list¶

y_list¶

rho_list¶

prev_pos = None¶

prev_forces = None¶

H0_scaling = 1.0¶

compute_step(pos: torch.Tensor, forces: torch.Tensor) → torch.Tensor¶

Compute the optimal step using L-BFGS approximation.

Parameters:

pos (torch.Tensor) – Current atomic positions (B×N×3 tensor).
forces (torch.Tensor) – Current forces (B×N×3 tensor).

Returns:

Optimal step vector (B×N×3 tensor).

Return type:

torch.Tensor

reset() → None¶: Reset the L-BFGS memory.

class agedi.diffusion.guidance.BatchedLBFGSStepSizer(batch_size: int, memory_size: int = 10, initial_step: float = 0.1)¶

Batched wrapper around LBFGSStepSizer for use with batched graphs.

Maintains one LBFGSStepSizer per graph in a batch and dispatches the step computation to the appropriate instance based on batch indices.

step_sizers¶

compute_step(pos: torch.Tensor, forces: torch.Tensor, batch_idx: torch.Tensor) → torch.Tensor¶

Compute steps for batched data.

Parameters:

pos (torch.Tensor) – Current atomic positions.
forces (torch.Tensor) – Current forces acting on the atoms.
batch_idx (torch.Tensor) – Index tensor mapping each atom to its graph in the batch.

Returns:

Combined step tensor with the same shape as pos.

Return type:

torch.Tensor

reset() → None¶: Reset the L-BFGS memory for all step-sizers in the batch.

agedi.diffusion.guidance.force_field_guidance_step(batch: agedi.data.AtomsGraph, regressor_model: torch.nn.Module, lbfgs_step_sizer: BatchedLBFGSStepSizer, scale: float, zeta: float = 3.0, max_step_size: float = 0.1) → agedi.data.AtomsGraph¶

Apply one force-field guidance step with batched L-BFGS step-size adaptation.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
regressor_model (torch.nn.Module) – The regressor model used to compute forces.
lbfgs_step_sizer (BatchedLBFGSStepSizer) – The L-BFGS step sizer (one per graph in the batch).
scale (float) – Base scale of the force field guidance.
zeta (float, optional) – Exponent for the time-dependent weight (1 - t)**zeta.
max_step_size (float, optional) – Maximum allowed step size magnitude. Default is 0.1.

Returns:

Updated batch after applying the guidance step.

Return type:

AtomsGraph

agedi.diffusion.guidance.post_diffusion_relaxation_step(batch: agedi.data.AtomsGraph, regressor_model: torch.nn.Module, lbfgs_step_sizer: BatchedLBFGSStepSizer | None, scale: float = 0.1) → agedi.data.AtomsGraph¶

Perform a pure force-based relaxation step after diffusion is complete.

Parameters:

batch (AtomsGraph) – A batch of AtomsGraph data.
regressor_model (torch.nn.Module) – The regressor model used to compute forces.
lbfgs_step_sizer (BatchedLBFGSStepSizer or None) – The L-BFGS step sizer. Initialised from batch if None.
scale (float, optional) – Step size scaling factor for relaxation.

Returns:

Updated batch after relaxation step.

Return type:

AtomsGraph