agedi¶

Submodules¶

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.

Functions¶

`create_dataset`(→ agedi.data.Dataset)	Create and setup an AGeDi Dataset from ASE Atoms objects.
`create_diffusion`(, sde, SDE] =, conditioning, ...)	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) → 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) → AtomsGraph¶

Create an empty graph.

Parameters:: cutoff (float) – The cutoff radius for the edges.
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_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)¶

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¶

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

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, **kwargs) → agedi.data.AtomsGraph¶

Initialise a single graph from noiser priors.

Parameters:

cutoff (float) – Cutoff radius for the neighbour list.
**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) → 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.

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, **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.
**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) → 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 when no template is given.
pbc – optional; defaults to [True, True, True].

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.
corrector_step_size (float, optional) – Step size for each corrector pass. Defaults to 1e-3.

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¶

agedi.create_dataset(data: Sequence[ase.Atoms], cutoff: float = 6.0, 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) → 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 = 6.0, feature_size: int = 64, n_blocks: int = 4, n_rbf: int = 30, noisers: Sequence[str | Noiser] = ('CellPositions',), sde: str | SDE = 've', 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) → 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

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¶

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