agedi.api.diffusion

Diffusion model creation and loading.

Functions

create_diffusion(, sde, SDE] =, conditioning, ...)

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

load_diffusion(→ Agedi)

Load a trained diffusion model from an AGeDi log directory.

Module Contents

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