agedi

Submodules

• agedi.api
• agedi.cli
• agedi.data
• agedi.diffusion
• agedi.functional
• agedi.models
• agedi.utils

Classes

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

Functions

create_dataset(→ agedi.data.Dataset)	Create and setup an AGeDi Dataset from ASE Atoms objects.
create_diffusion(, sde, SDE, None] = None, ...)	Create a diffusion model for script-based training and sampling.
create_trainer(→ lightning.Trainer)	Create a Lightning trainer configured for AGeDi.
load_diffusion(→ Agedi)	Load a trained diffusion model from an AGeDi log directory.
predict(→ List[ase.Atoms])	Predict energies and forces for input structures using a trained force-field.
register_model(→ None)	Register a custom score model backbone factory under name.
sample(→ Union[List[agedi.data.AtomsGraph], ...)	Sample structures from a trained diffusion model.
train(→ lightning.Trainer)	Train a diffusion model and return the trainer used.
train_from_atoms(*args, **kwargs)
train_from_config(*args, **kwargs)

Package Contents

class agedi.AtomsGraph

Bases: torch_geometric.data.Data

Atomistic Graph Class

Class defining a graph with atoms as nodes and edges formed between all atoms within a finite cutoff radius.

Parameters:

• pos (torch.Tensor) – The positions of the atoms with shape (n_atoms, 3).

• x (torch.Tensor) – The node features i.e atomic types of the graph with shape (n_nodes, 1).

• edge_index (torch.Tensor) – The edge index tensor of the graph with shape (2, n_edges).

• edge_attr (torch.Tensor) – The edge attributes of the graph with shape (n_edges, n_edge_features).

• y (Optional[torch.Tensor]) – The target tensor of the graph with shape (n_targets,).

• representation (Optional[Representation]) – The representation of the atoms in the graph.

• confinement (Optional[torch.Tensor]) – z-directional confinement of the atoms with shape (1,2).

• kwargs (Dict[str, torch.Tensor])

classmethod from_atoms(atoms: ase.Atoms, cutoff: float = 6.0, dtype: torch.dtype = torch.float, initialize_mask: bool | None = None, confinement: Tuple[float, float] | None = None, canonical_cell: bool = False, fully_connected: bool = False) → AtomsGraph

Create a graph from an ASE Atoms object.

Parameters:

• atoms (Atoms) – The ASE Atoms object.

• cutoff (float) – The cutoff radius for the edges.

• dtype (torch.dtype) – The data type of the tensors.

• initialize_mask (Optional[bool]) – Whether to initialize the mask tensor. When None (the default), the mask is initialised only when confinement is not provided (i.e. initialize_mask defaults to False for template / confinement graphs).

• confinement (Optional[Tuple[float, float]]) – Optional z-directional confinement bounds (z_min, z_max) to attach to the graph. When provided, a confinement tensor of shape (1, 2) is stored on the graph. When None (the default), no confinement attribute is added.

• canonical_cell (bool) – When True, the cell is stored in canonical lower-triangular form. If the input cell is not already canonical, Cartesian positions are recomputed to preserve fractional coordinates and a warning is raised. Set to False (the default) to store the cell exactly as provided by ASE (no rotation or recomputation is performed).

Returns:

graph – The graph object.

Return type:

AtomsGraph

classmethod empty(cutoff: float = 6.0, fully_connected: bool = False) → AtomsGraph

Create an empty graph.

Parameters:

• cutoff (float) – The cutoff radius for the edges.

• fully_connected (bool, optional) – When True the graph will be rebuilt as a fully connected graph (all atom pairs, no self-loops, zero shift vectors) every time update_graph() is called. Suitable for gas-phase molecules and clusters where a finite cutoff misses pairs when atoms spread during the reverse diffusion process. Defaults to False.

Returns:

graph – The graph object.

Return type:

AtomsGraph

add_batch_attr(key: str, value: torch.Tensor, type: str = 'node') → None

Add a batch attribute to the graph.

Parameters:

• key (str) – The key of the attribute.

• value (torch.Tensor) – The value of the attribute.

• type (str) – The type of the attribute. Can be either “node” or “graph”

Return type:

None

to_atoms() → ase.Atoms

Convert the graph to an ASE Atoms object.

Only works on unbatched graphs.

Returns:

atoms – The atoms object.

Return type:

ase.Atoms

_get_scalar_attr(key: str) → float | None

prepare_for_compile(cutoff: float) → None

Pre-allocate neighbor-list buffers for torch.compile compatibility.

Estimates the maximum number of neighbors per atom using estimate_max_neighbors() and the cell-list dimensions using estimate_cell_list_sizes(), then allocates the cell list and all output buffers with fixed shapes. Fixed shapes are required for torch.compile to trace the reverse diffusion step once without retracing on subsequent iterations.

Must be called on a Batch before the first update_graph() call.

Requires the nvalchemiops package.

Parameters:

cutoff (float) – Neighbor-list cutoff radius (Å).

Raises:

• RuntimeError – When nvalchemiops is not installed.

• TypeError – When called on an unbatched AtomsGraph instead of a Batch.

static _cell_list_to_graph(neighbor_matrix: torch.Tensor, neighbor_shifts: torch.Tensor, cell: torch.Tensor, dtype: torch.dtype, batch_idx: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]

Convert cell-list query output to (edge_index, shift_vectors).

update_graph() → bool

Update the graph with new edges

This should be called after changing any of the positions or cell.

Returns:

rebuilt – True when the neighbor list was fully recomputed.

Return type:

bool

static _make_graph_matscipy(positions: torch.Tensor, cell: torch.Tensor, cutoff: float, pbc: torch.Tensor, dtype: torch.dtype | None = None, batch_idx: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]

static make_fully_connected_graph(positions: torch.Tensor, dtype: torch.dtype | None = None, batch_idx: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]

Build a fully connected graph: every atom is connected to every other.

No self-loops are included. All shift vectors are zero (non-periodic). This is the correct topology for gas-phase molecules and clusters where the finite cutoff of a standard neighbour list would miss long-range pairs when atoms spread apart during the reverse diffusion process.

Parameters:

• positions (torch.Tensor, shape (n_atoms, 3))

• dtype (torch.dtype, optional) – Data type of the shift-vector output. Defaults to positions.dtype.

• batch_idx (torch.Tensor of shape (n_atoms,), optional) – Graph-membership index for batched graphs. When None all atoms are treated as belonging to a single graph.

Returns:

• edge_index (torch.Tensor, shape (2, n_edges))

• shift_vectors (torch.Tensor, shape (n_edges, 3)) – All zeros (no periodic images).

static make_graph(positions: torch.Tensor, cell: torch.Tensor, cutoff: float, pbc: torch.Tensor, dtype: torch.dtype = None, batch_idx: torch.Tensor | None = None) → Tuple[torch.Tensor, torch.Tensor]

Create the graph-edges from the positions and cell.

Parameters:

• positions (torch.Tensor) – The positions of the atoms.

• cell (torch.Tensor) – The cell of the system.

• cutoff (float) – The cutoff radius for the edges.

• pbc (torch.Tensor) – The periodic boundary conditions.

• dtype (torch.dtype) – The data type of the output.

Returns:

• edge_index (torch.Tensor) – The edge index tensor.

• shift_vectors (torch.Tensor) – The shift vectors tensor.

clear_graph() → None

Clear the graph removing all edges

Return type:

None

__len__() → int

Return the number of atoms in the graph.

Returns:

n_atoms – The number of atoms in the graph.

Return type:

int

property cell: torch.Tensor

Return the canonical cell matrix of the graph.

Returns:

cell – The cell matrix of shape (3, 3).

Return type:

torch.Tensor

property frac: torch.Tensor

Return the fractional coordinates of the positions

Returns:

frac – The fractional coordinates of the atoms.

Return type:

torch.Tensor

frac_to_pos(f: torch.Tensor) → torch.Tensor

Fraction -> Cartesian coordinates.

Convert fractional coordinates to cartesian coordinates.

Parameters:

f (torch.Tensor) – The fractional coordinates.

Returns:

r – The cartesian coordinates.

Return type:

torch.Tensor

pos_to_frac(r: torch.Tensor) → torch.Tensor

Cartesian -> Fractional coordinates.

Convert cartesian coordinates to fractional coordinates.

Parameters:

r (torch.Tensor) – The cartesian coordinates.

Returns:

f – The fractional coordinates.

Return type:

torch.Tensor

property positions_mask: torch.Tensor

Return the mask of the positions that are fixed.

True for fixed atom-positions and else false.

Returns:

mask – The mask of the positions that are fixed.

Return type:

torch.Tensor

property time: torch.Tensor

Return the time of the graph.

Returns:

time – The time of the graph.

Return type:

torch.Tensor

property representation: Representation | None

Return the representation of the graph.

Returns:

representation – The representation of the graph, or None if not set.

Return type:

Optional[Representation]

wrap_positions() → None

Wrap the positions of the atoms to the unit cell.

Return type:

None

apply_mask(x: torch.Tensor, val: float = 0.0) → torch.Tensor

Apply the mask to the tensor x.

Parameters:

• x (torch.Tensor) – The tensor to apply the mask to.

• val (float) – The value to set the masked values to.

Returns:

x – The tensor with the mask applied.

Return type:

torch.Tensor

property confinement: torch.Tensor

Return the confinement of the graph.

Returns:

confinement – The confinement of the graph.

Return type:

torch.Tensor

property cellpar: torch.Tensor

Return the cell parameters of the graph.

static _is_lower_triangular(cell: torch.Tensor) → bool

Return True if cell is in canonical lower-triangular form.

A cell matrix is considered canonical when the three strictly upper-triangular entries (cell[0,1], cell[0,2], cell[1,2]) are all zero (within a tight floating-point tolerance of 1e-10).

Parameters:

cell (torch.Tensor) – The cell matrix.

Returns:

True if the cell is already lower-triangular.

Return type:

bool

static cell_to_vectors(cell: torch.Tensor) → torch.Tensor

Convert cell matrix to cell parameters.

Parameters:

cell (torch.Tensor) – The cell matrix of shape (N, 3) or (N, 3, 3).

Returns:

The cell parameters of shape (N, 6).

Return type:

torch.Tensor

static vector_to_cell(cellpar: torch.Tensor) → torch.Tensor

Convert cell parameters to cell matrix.

Parameters:

cellpar (torch.Tensor) – The cell parameters of shape (N, 6).

Returns:

The cell matrix of shape (N, 3, 3) where each row is a lattice vector.

Return type:

torch.Tensor

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

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

Parameters:

• name (str) – Alias used to look up this sampler, e.g. "heun".

• factory (callable) – factory(score_fn, noisers, **kwargs) -> Sampler.

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))

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

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

1. Evaluate score model.

2. Apply each noiser’s denoising update in reverse order.

3. Wrap positions and rebuild the neighbour list.

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

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.

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

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

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.

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

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

agedi.create_dataset(data: Sequence[ase.Atoms], cutoff: float | None = None, batch_size: int = 64, train_split: float | int = 0.9, val_split: float | int = 0.1, mask: str = 'none', confinement: Tuple[float, float] | None = None, conditioning: str = 'none', conditioning_type: str = 'scalar', repeat: int | None = None, canonical_cell: bool = False, regressor_data: Sequence[ase.Atoms] | None = None, properties: List[Dict] | None = None, fully_connected: bool = False) → agedi.data.Dataset

Create and setup an AGeDi Dataset from ASE Atoms objects.

Parameters:

• data (Sequence[Atoms]) – ASE Atoms objects to add to the dataset.

• cutoff (float, optional) – Neighbour-list cutoff radius in Ångström.

• batch_size (int, optional) – Mini-batch size used during training/validation.

• train_split (Union[float, int], optional) – Fraction or absolute number of samples for the training split.

• val_split (Union[float, int], optional) – Fraction or absolute number of samples for the validation split.

• mask (str, optional) – Atom-mask method (e.g. "MaskFixed" or "none").

• confinement (Tuple[float, float], optional) – Z-axis confinement bounds (z_min, z_max).

• conditioning (str, optional) – Name of the per-structure property to use as a conditioning signal. The value is read from atoms.info[conditioning] or the corresponding atoms.get_<conditioning>() method. Ignored when set to "none" (default).

• conditioning_type (str, optional) – "scalar" (default) or "node"; controls how the conditioning property is broadcast onto the graph.

• repeat (int, optional) – When given, augment the dataset by repeating each structure up to repeat times along the first two cell vectors.

• canonical_cell (bool, optional) – Store cells in canonical lower-triangular form.

• regressor_data (Sequence[Atoms], optional) – Additional ASE Atoms objects used to train a regressor head.

• properties (List[Dict], optional) – Per-structure property dictionaries; must contain exactly one entry per element in data. Each dictionary is merged into the corresponding graph object via setattr, matching the layout accepted by add_atoms_data(). Keys already produced by the conditioning logic are overwritten by values in properties when both are present.

Returns:

A fully set-up Dataset ready for training.

Return type:

Dataset

agedi.create_diffusion(model: str = 'PaiNN', cutoff: float | None = None, feature_size: int = 64, n_blocks: int = 4, n_rbf: int = 30, noisers: Sequence[str | Noiser] = ('CellPositions',), sde: str | SDE | None = None, conditioning: str = 'none', conditioning_type: str = 'scalar', confinement: Tuple[float, float] | None = None, force_field: bool = False, lr: float = 0.0001, lr_factor: float = 0.95, lr_patience: int = 100, weight_decay: float = 0.0, eps: float = 1e-05, guidance_weight: float = -1.0, device: str | torch.device | None = None, type_map: List[int] | None = None, prediction_type: str = 'score', sampler: str = 'em', loss_weighting: str = 'uniform', fully_connected: bool = False) → agedi.Agedi

Create a diffusion model for script-based training and sampling.

Parameters:

• model (str, optional) – GNN backbone architecture. The name is looked up in the model registry; use register_model() to add custom backends. The built-in default is "PaiNN" (SchNetPack PaiNN).

• cutoff (float, optional) – Neighbour-list cutoff radius in Å. Defaults to 6.0.

• feature_size (int, optional) – Embedding / feature dimension. Defaults to 64.

• n_blocks (int, optional) – Number of interaction blocks. Defaults to 4.

• n_rbf (int, optional) – Number of radial basis functions. Defaults to 30.

• noisers (Sequence[str or Noiser], optional) –

  Noiser identifiers or instances to include. Defaults to ("CellPositions",). Recognised string identifiers (CamelCase preferred; snake_case aliases also accepted for backwards compatibility):

  • "Positions" / "positions" – Positions (StandardNormal prior + Normal, for gas-phase clusters).

  • "CellPositions" / "cell_positions" – CellPositions (UniformCell prior + Normal, for periodic bulk/surface systems).

  • "ConfinedCellPositions" / "confined_cell_positions" – ConfinedCellPositions (UniformCellConfined prior + TruncatedNormal, for Z-confined systems).

  • "Types" / "types" – Types.

• sde (str or SDE, optional) – SDE for position noisers. Short aliases: "ve" (default), "vp". Pass an instantiated SDE for full control.

• conditioning (str, optional) – Property to condition on, or "none" for time-only conditioning. Defaults to "none".

• conditioning_type (str, optional) – Type of the conditioning module: "scalar" or "integer". Defaults to "scalar".

• confinement (Tuple[float, float], optional) – Z-direction confinement bounds (z_min, z_max) in Å.

• force_field (bool, optional) – When True, attach a diffusion.regressor_model. The heads shares the same representation and translator as the score model so that atomic embeddings are learned jointly. It is trained whenever the training batch contains per-atom forces and total energies (i.e. the ASE training structures have DFT (or other) energy and forces). The trained forces head enables force-field guided sampling via ForcefieldGuidanceConfig. Defaults to False.

• lr (float, optional) – Learning rate. Defaults to 1e-4.

• lr_factor (float, optional) – LR-scheduler reduction factor. Defaults to 0.95.

• lr_patience (int, optional) – LR-scheduler patience (epochs). Defaults to 100.

• weight_decay (float, optional) – Optimizer weight-decay. Defaults to 0.0.

• eps (float, optional) – Minimum diffusion time. Defaults to 1e-5.

• guidance_weight (float, optional) – Classifier-free guidance weight. Defaults to -1.0 (disabled).

• device (str or torch.device, optional) – Target compute device. When None CUDA is used if available, otherwise CPU.

• type_map (List[int], optional) – Compact type map for the Types noiser. type_map[0] must be 0 (absorbing state) and type_map[i] is the atomic number for compact index i. When provided, the Types noiser and the TypesScore head use a reduced vocabulary of size len(type_map) instead of the default 100. Auto-populated by train_from_atoms() when a "Types" noiser is requested.

Returns:

A freshly initialised Agedi model.

Return type:

Agedi

agedi.create_trainer(*, epochs: int = -1, max_time: int | Dict | datetime.timedelta | None = 24, accelerator: str = 'auto', devices: int = 1, logger: str = 'tensorboard', log_dir: str = 'logs', project: str = 'agedi', name: str = 'agedi', log_interval: int = 10, gradient_clip_val: float = 10.0, progress_bar: bool = False, print_epoch_interval: int = 10, log_grad_norm: bool = True, repeat: int | None = None, repeat_epoch: int | None = None, hparams: Dict | None = None, extra_callbacks: List[lightning.pytorch.callbacks.Callback] | None = None) → lightning.Trainer

Create a Lightning trainer configured for AGeDi.

Parameters:

• epochs – Maximum number of training epochs (-1 = unlimited).

• max_time –

  Wall-clock time limit for training. Accepts:

  • int – number of hours (e.g. 24 ≡ 24 hours).

  • dict – Lightning-style mapping, e.g. {"days": 0, "hours": 12, "minutes": 30, "seconds": 0}.

  • datetime.timedelta – a Python timedelta object.

  • None – no time limit.

• accelerator – Hardware accelerator to use (e.g. "auto", "gpu", "cpu"). Default: "auto".

• devices – Number of devices to train on. Default: 1.

• logger – Logging backend: "tensorboard" (default) or "wandb".

• log_dir – Root directory for logs and checkpoints. Default: "logs".

• project – WandB project name (only used when logger="wandb").

• name – Experiment display name used by TensorBoard and WandB as the run sub-directory / run name. Default: "agedi".

• log_interval – How often (in steps) to log metrics. Default: 10.

• gradient_clip_val – Maximum gradient norm for gradient clipping. Default: 10.0.

• progress_bar – Whether to show a Lightning progress bar. Default: False.

• print_epoch_interval – Print a one-line training summary to stdout every this many epochs. Set to 0 to disable. Default: 10.

• log_grad_norm – Whether to log the total gradient norm during training. Disable for large models where the per-step overhead is undesirable. Default: True.

• repeat – Number of repetition levels for cell-repeat data augmentation. Must be set together with repeat_epoch. When None (default), no repetition augmentation is applied.

• repeat_epoch – How many epochs between repetition-level increases. Required when repeat is set.

• hparams – Hyperparameters dict logged to hparams.yaml via HParamsMetricLogger. When None (default), no extra hyperparameter logging is performed.

• extra_callbacks – Extra Lightning callbacks to append to the default callback list. When None (default) only the built-in callbacks are used.

Returns:

A configured Trainer ready to call trainer.fit(diffusion, dataset).

Return type:

lightning.Trainer

agedi.load_diffusion(path: str | pathlib.Path, checkpoint: str | pathlib.Path | None = None, device: str | torch.device | None = None) → Agedi

Load a trained diffusion model from an AGeDi log directory.

The model architecture is fully reconstructed from the Hydra-compatible diffusion config stored in hparams.yaml, so no additional parameters are needed.

Parameters:

• path – Path to the AGeDi log / model directory (or directly to the hparams.yaml file).

• checkpoint – Path to a specific checkpoint file. When None the latest checkpoint (checkpoints/last_model.ckpt) is loaded automatically.

• device – Device to load the model onto. When None CUDA is used if available, otherwise CPU.

agedi.predict(diffusion: Agedi, structures: Sequence[ase.Atoms], *, batch_size: int = 64, cutoff: float | None = None) → List[ase.Atoms]

Predict energies and forces for input structures using a trained force-field.

The model must have been trained with force_field=True (i.e. it must have a regressor_model attached). The predicted energy and forces are attached to the returned Atoms objects via an SinglePointCalculator.

Parameters:

• diffusion – A trained Agedi model with a force-field regressor (trained with --force_field).

• structures – Input ASE Atoms objects to run predictions on.

• batch_size – Number of structures per inference batch. Defaults to 64.

• cutoff – Neighbour-list cutoff in Å. When None (default), the cutoff is read from the model’s representation automatically.

Returns:

The input structures with a SinglePointCalculator attached containing the predicted energy and/or forces.

Return type:

List[Atoms]

Raises:

ValueError – If the model does not have a force-field regressor.

agedi.register_model(name: str, factory: Callable) → None

Register a custom score model backbone factory under name.

The factory is called with the keyword arguments cutoff, heads, feature_size, n_blocks, head_dim, and n_rbf and must return a 3-tuple (translator, representation, List[Head]).

Registered models can be selected by passing model=name to create_diffusion().

Parameters:

• name (str) – Alias used to select this backend (e.g. "PaiNN").

• factory (Callable) –

  Factory function with signature:

  factory(cutoff, heads, feature_size, n_blocks, head_dim, n_rbf)
      -> Tuple[Translator, nn.Module, List[Head]]
  

Examples

from agedi.functional import register_model

def my_factory(cutoff, heads, feature_size, n_blocks, head_dim, n_rbf):
    ...
    return translator, representation, head_list

register_model("MyModel", my_factory)

agedi.sample(diffusion: Agedi, *, n_samples: int, 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, template: agedi.data.AtomsGraph | ase.Atoms | None = None, confinement: Tuple[float, float] | None = None, compile: bool = False, steps: int = 500, eps: float = 0.001, batch_size: int = 64, ff_guidance: agedi.diffusion.ForcefieldGuidanceConfig | None = None, property: Dict[str, float] | None = None, progress_bar: bool = False, save_trajectory: bool = False, print_timings: bool = False, as_atoms: bool = True, sampler=None, sampler_kwargs=None) → List[agedi.data.AtomsGraph] | List[ase.Atoms] | List[List[agedi.data.AtomsGraph]] | List[List[ase.Atoms]]

Sample structures from a trained diffusion model.

Parameters:

• diffusion – A trained Agedi model.

• n_samples – Number of structures to generate.

• n_atoms – Number of atoms per structure. Automatically determined from formula if provided, or from the length of atomic_numbers when n_atoms is not explicitly given.

• atomic_numbers – Atomic numbers of the generated atoms. Not required when the model has a types-noiser or when formula is provided.

• formula – Chemical formula (e.g. "H2O"). Used to derive n_atoms and atomic_numbers when they are not provided explicitly.

• positions – Fixed positions of the atoms (shape (n_atoms, 3)). Required when no positions-noiser is configured (type-only diffusion). Positions will not be modified during sampling.

• cell – Unit-cell matrix (3×3 array or flat length-9 array). Not required when template is provided (the template’s cell is used instead).

• pbc – Periodic boundary conditions as a length-3 boolean array (e.g. [True, True, False]). When template is provided its pbc is used unless this argument is given explicitly. Defaults to [True, True, True] (fully periodic) when neither template nor pbc is supplied.

• template – Template structure. May be an AtomsGraph or an ASE Atoms object; the latter is automatically converted to an AtomsGraph (with confinement applied when provided). When given, cell and pbc are taken from the template unless explicitly provided.

• ff_guidance – Force-field guidance configuration. When None (default) a ForcefieldGuidanceConfig with default values is used (i.e. guidance is disabled).

• compile – When True, use torch.compile on the reverse diffusion step for faster sampling. Before the sampling loop starts, the maximum number of neighbors and cell-list dimensions are estimated automatically via NVIDIA nvalchemiops (estimate_max_neighbors and estimate_cell_list_sizes), and all neighbor-list buffers are pre-allocated with fixed shapes. Requires NVIDIA nvalchemiops. Defaults to False.

• print_timings – When True, print a per-stage timing breakdown at the end of each sampling batch (graph init, score model, denoise, neighbor list, etc.). Defaults to False.

agedi.train(diffusion: Agedi, dataset: agedi.data.Dataset, trainer: lightning.Trainer | None = None, ckpt_path: str | pathlib.Path | None = None, **trainer_kwargs) → lightning.Trainer

Train a diffusion model and return the trainer used.

Parameters:

• diffusion – The diffusion model to train.

• dataset – The dataset to train on.

• trainer – A pre-configured Lightning Trainer. When None a new trainer is created from trainer_kwargs.

• ckpt_path – Path to a Lightning checkpoint (.ckpt) to resume training from. When provided the full training state (model weights, optimiser, LR-scheduler, and epoch counter) is restored before fitting. Equivalent to passing ckpt_path to trainer.fit().

• **trainer_kwargs – Additional keyword arguments forwarded to create_trainer() when trainer is None.

agedi.train_from_atoms(*args, **kwargs)

agedi.train_from_config(*args, **kwargs)