agedi.diffusion.noisers.base ============================ .. py:module:: agedi.diffusion.noisers.base Classes ------- .. autoapisummary:: agedi.diffusion.noisers.base.Noiser Module Contents --------------- .. py:class:: Noiser(distribution: agedi.diffusion.distributions.Distribution, prior: agedi.diffusion.distributions.Distribution, loss_scaling: float = 1.0, key: Optional[str] = None, **kwargs) Bases: :py:obj:`abc.ABC`, :py:obj:`torch.nn.Module` Noiser Base class Impments a noiser that can noise and denoise a atomistic structure attribute. :param distribution: The distribution to be used for the noising. :type distribution: Distribution :param prior: The prior to be used for the denoising. :type prior: Distribution :param loss_scaling: Scaling factor applied to this noiser's loss contribution. :type loss_scaling: float :param key: Override the class-level ``_key`` for the attribute to be noised and denoised. Useful for reusing a noiser class on a different attribute without subclassing purely to change ``_key``. :type key: str, optional :rtype: Noiser .. py:attribute:: _key :type: str .. py:attribute:: _registry :type: ClassVar[Dict[str, Callable[..., Noiser]]] .. py:attribute:: distribution .. py:attribute:: prior .. py:attribute:: loss_scaling :value: 1.0 .. py:method:: register(name: str, factory: Callable[..., Noiser]) -> None :classmethod: Register a noiser factory callable under *name*. The factory is called with ``sde`` as a keyword argument containing the resolved :class:`~agedi.diffusion.sdes.SDE` instance. Noisers that do not use an SDE can safely ignore it via ``**kwargs``. :param name: Alias string used to reference the noiser in :func:`~agedi.functional.create_diffusion`. :type name: str :param factory: A callable that accepts ``sde`` as a keyword argument and returns a :class:`Noiser` instance. :type factory: Callable .. rubric:: Examples Register a custom noiser so it can be referenced by its alias:: from agedi.diffusion.noisers import Noiser class MyNoiser(Noiser): ... Noiser.register("my_noiser", lambda sde: MyNoiser(sde=sde)) .. py:method:: get_hparams() -> Dict Return hyperparameters sufficient to reconstruct this noiser. Returns a dictionary with a ``_target_`` key (the fully-qualified class name) plus ``distribution``, ``prior``, and ``loss_scaling`` entries taken from the base class. Subclasses should call ``super().get_hparams()`` and merge in their own constructor parameters. :returns: Hyperparameter dictionary. :rtype: dict .. py:property:: key :type: str The key of the attribute to be noised and denoised. .. py:method:: _noise(batch: agedi.data.AtomsGraph) -> agedi.data.AtomsGraph :abstractmethod: Noises the attribute of the atomistic structure. Must be implemented by the subclass. :param batch: The atomistic structure (or batch hereof) to be noised. :type batch: AtomsGraph :returns: The noised atomistic structure (or bach hereof). :rtype: AtomsGraph .. py:method:: _denoise(batch: agedi.data.AtomsGraph, delta_t: float, last: bool) -> agedi.data.AtomsGraph :abstractmethod: Denoises the attribute of the atomistic structure. Must be implemented by the subclass. :param batch: The atomistic structure (or batch hereof) to be denoised. :type batch: AtomsGraph :param delta_t: The time step to be used for the denoising. :type delta_t: float :returns: The denoised atomistic structure (or bach hereof). :rtype: AtomsGraph .. py:method:: _loss(batch: agedi.data.AtomsGraph) -> float :abstractmethod: Computes the training loss. Must be implemented by the subclass. :param batch: The atomistic structure (or batch hereof) to be noised and denoised. :type batch: AtomsGraph :returns: The loss of the noised and denoised atomistic structure. :rtype: float .. py:method:: noise(batch: agedi.data.AtomsGraph) -> agedi.data.AtomsGraph Noises the attribute of the atomistic structure. :param batch: The atomistic structure (or batch hereof) to be noised. :type batch: AtomsGraph :returns: The noised atomistic structure (or bach hereof). :rtype: AtomsGraph .. py:method:: denoise(batch: agedi.data.AtomsGraph, delta_t: float, last: bool) -> agedi.data.AtomsGraph Denoises the attribute of the atomistic structure. :param batch: The atomistic structure (or batch hereof) to be denoised. :type batch: AtomsGraph :param delta_t: The time step to be used for the denoising. :type delta_t: float :param last: If the denoising is the last step of the denoising. :type last: bool :returns: The denoised atomistic structure (or bach hereof). :rtype: AtomsGraph .. py:method:: loss(batch: agedi.data.AtomsGraph) -> float Compute the training loss. :param batch: The atomistic structure (or batch hereof) to be noised and denoised. :type batch: AtomsGraph :returns: The loss of the noised and denoised atomistic structure. :rtype: float .. py:method:: langevin_step(batch: agedi.data.AtomsGraph, step_size: Union[float, torch.Tensor] = 0.01) -> agedi.data.AtomsGraph Perform a Langevin corrector step at the current (constant) time. Applies a small score-corrected Langevin update without advancing the diffusion time. The score must already be stored in ``batch[key + "_score"]`` (i.e. the score model must have been called before invoking this method). The default implementation delegates to :meth:`_denoise` with ``last=False`` and a fixed *step_size*. Subclasses may override this for a more specialised corrector. :param batch: The atomistic structure (or batch hereof) to be corrected. :type batch: AtomsGraph :param step_size: Size of the Langevin corrector step. Passing a pre-created :class:`torch.Tensor` avoids repeated tensor allocation when this method is called in a tight loop. Defaults to ``0.01``. :type step_size: float or torch.Tensor, optional :returns: The corrected atomistic structure. :rtype: AtomsGraph .. py:method:: initialize_graph(batch: agedi.data.AtomsGraph) -> None Initializes the graph with the prior distribution. Can be overwritten by subclasses for specific initializations. :param batch: The atomistic structure (or batch hereof) to be noised and denoised. :type batch: AtomsGraph