agedi.api.diffusion =================== .. py:module:: agedi.api.diffusion .. autoapi-nested-parse:: Diffusion model creation and loading. Functions --------- .. autoapisummary:: agedi.api.diffusion.create_diffusion agedi.api.diffusion.load_diffusion Module Contents --------------- .. py:function:: create_diffusion(model: str = 'PaiNN', cutoff: float = 6.0, feature_size: int = 64, n_blocks: int = 4, n_rbf: int = 30, noisers: Sequence[Union[str, Noiser]] = ('CellPositions', ), sde: Union[str, SDE] = 've', conditioning: str = 'none', conditioning_type: str = 'scalar', confinement: Optional[Tuple[float, float]] = 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: Optional[Union[str, torch.device]] = None, type_map: Optional[List[int]] = None) -> agedi.Agedi Create a diffusion model for script-based training and sampling. :param model: GNN backbone architecture. The name is looked up in the model registry; use :func:`register_model` to add custom backends. The built-in default is ``"PaiNN"`` (SchNetPack PaiNN). :type model: str, optional :param cutoff: Neighbour-list cutoff radius in Å. Defaults to ``6.0``. :type cutoff: float, optional :param feature_size: Embedding / feature dimension. Defaults to ``64``. :type feature_size: int, optional :param n_blocks: Number of interaction blocks. Defaults to ``4``. :type n_blocks: int, optional :param n_rbf: Number of radial basis functions. Defaults to ``30``. :type n_rbf: int, optional :param noisers: Noiser identifiers or instances to include. Defaults to ``("CellPositions",)``. Recognised string identifiers (CamelCase preferred; snake_case aliases also accepted for backwards compatibility): * ``"Positions"`` / ``"positions"`` – :class:`~agedi.diffusion.noisers.Positions` (StandardNormal prior + Normal, for gas-phase clusters). * ``"CellPositions"`` / ``"cell_positions"`` – :class:`~agedi.diffusion.noisers.CellPositions` (UniformCell prior + Normal, for periodic bulk/surface systems). * ``"ConfinedCellPositions"`` / ``"confined_cell_positions"`` – :class:`~agedi.diffusion.noisers.ConfinedCellPositions` (UniformCellConfined prior + TruncatedNormal, for Z-confined systems). * ``"Types"`` / ``"types"`` – :class:`~agedi.diffusion.noisers.Types`. :type noisers: Sequence[str or Noiser], optional :param sde: SDE for position noisers. Short aliases: ``"ve"`` (default), ``"vp"``. Pass an instantiated :class:`~agedi.diffusion.sdes.SDE` for full control. :type sde: str or SDE, optional :param conditioning: Property to condition on, or ``"none"`` for time-only conditioning. Defaults to ``"none"``. :type conditioning: str, optional :param conditioning_type: Type of the conditioning module: ``"scalar"`` or ``"integer"``. Defaults to ``"scalar"``. :type conditioning_type: str, optional :param confinement: Z-direction confinement bounds ``(z_min, z_max)`` in Å. :type confinement: Tuple[float, float], optional :param force_field: 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 :class:`~agedi.diffusion.ForcefieldGuidanceConfig`. Defaults to ``False``. :type force_field: bool, optional :param lr: Learning rate. Defaults to ``1e-4``. :type lr: float, optional :param lr_factor: LR-scheduler reduction factor. Defaults to ``0.95``. :type lr_factor: float, optional :param lr_patience: LR-scheduler patience (epochs). Defaults to ``100``. :type lr_patience: int, optional :param weight_decay: Optimizer weight-decay. Defaults to ``0.0``. :type weight_decay: float, optional :param eps: Minimum diffusion time. Defaults to ``1e-5``. :type eps: float, optional :param guidance_weight: Classifier-free guidance weight. Defaults to ``-1.0`` (disabled). :type guidance_weight: float, optional :param device: Target compute device. When ``None`` CUDA is used if available, otherwise CPU. :type device: str or torch.device, optional :param type_map: Compact type map for the :class:`~agedi.diffusion.noisers.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 :func:`train_from_atoms` when a ``"Types"`` noiser is requested. :type type_map: List[int], optional :returns: A freshly initialised :class:`~agedi.Agedi` model. :rtype: Agedi .. py:function:: load_diffusion(path: Union[str, pathlib.Path], checkpoint: Optional[Union[str, pathlib.Path]] = None, device: Optional[Union[str, torch.device]] = 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. :param path: Path to the AGeDi log / model directory (or directly to the ``hparams.yaml`` file). :param checkpoint: Path to a specific checkpoint file. When ``None`` the latest checkpoint (``checkpoints/last_model.ckpt``) is loaded automatically. :param device: Device to load the model onto. When ``None`` CUDA is used if available, otherwise CPU.