agedi.diffusion.noisers ======================= .. py:module:: agedi.diffusion.noisers Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/agedi/diffusion/noisers/base/index /autoapi/agedi/diffusion/noisers/pos/index /autoapi/agedi/diffusion/noisers/sde/index /autoapi/agedi/diffusion/noisers/types/index Attributes ---------- .. autoapisummary:: agedi.diffusion.noisers.TypesNoiser Classes ------- .. autoapisummary:: agedi.diffusion.noisers.Noiser agedi.diffusion.noisers.SDENoiser agedi.diffusion.noisers.PositionsNoiser agedi.diffusion.noisers.Positions agedi.diffusion.noisers.CellPositions agedi.diffusion.noisers.ConfinedCellPositions agedi.diffusion.noisers.Types Package 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 .. py:class:: SDENoiser(sde_class: agedi.diffusion.sdes.SDE, sde_kwargs: Optional[Dict], distribution: agedi.diffusion.distributions.Distribution, prior: agedi.diffusion.distributions.Distribution, sde: Optional[agedi.diffusion.sdes.SDE] = None, **kwargs) Bases: :py:obj:`agedi.diffusion.noisers.Noiser`, :py:obj:`abc.ABC` Implements a SDE base class that can be inherited by other classes. :param sde_class: The class of the SDE to be used for the noising. :type sde_class: SDE :param sde_kwargs: The keyword arguments to be passed to the SDE class. :type sde_kwargs: Dict :param distribution: The distribution to be used for the noise. :type distribution: Distribution :param prior: The prior distribution to be used for the noise. :type prior: Distribution :param sde: An already-instantiated SDE object. When provided, *sde_class* and *sde_kwargs* are ignored. :type sde: SDE, optional :param key: The key to be used for the noising. :type key: str :param \*\*kwargs: Additional keyword arguments to be passed to the Noiser class. :returns: The noiser for the atoms positions in Cartesian coordinates. :rtype: Noiser .. py:attribute:: _key :value: None .. py:method:: get_hparams() -> Dict Return hyperparameters for this SDE noiser. .. py:method:: postprocess_score(score: torch.Tensor) -> torch.Tensor :abstractmethod: Post-process the predicted score before computing the loss. :param score: Raw predicted score tensor. :type score: torch.Tensor :returns: Post-processed score tensor. :rtype: torch.Tensor .. py:method:: postprocess_noise(noise: torch.Tensor) -> torch.Tensor :abstractmethod: Post-process the noise tensor before computing the loss. :param noise: Raw noise tensor. :type noise: torch.Tensor :returns: Post-processed noise tensor. :rtype: torch.Tensor .. py:method:: _noise(batch: agedi.data.AtomsGraph) -> agedi.data.AtomsGraph Adds noise to the atomistic structure. Added noise is stored in the self.key+"_noise". :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 positions of the atomistic structure. The denoising follows the Euler-Maruyama scheme. ::math:: R_i+1 = R_i + \Delta t (f(R_i, t) + g(t)**2 * s(R_i, t)) + \sqrt{\Delta t} g(t) * w The used score is expected to be stored in the self.key+"_score". :param batch: The atomistic structure (or batch hereof) to be denoised. :type batch: AtomsGraph :param delta_t: The time step 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) -> torch.Tensor Compute the noiser loss. Computes the loss of the diffusion model SDE noiser Expects the total added noise to be stored in the self.key+"_noise", and the predicted score to be stored in the self.key+"_score". The loss is computed as ::math:: L = \sum_i ||\sigma_t w_i + \sigma_t^2 s(R_i)||^2 :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:class:: PositionsNoiser(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Optional[Dict] = None, distribution: agedi.diffusion.distributions.Distribution = Normal(), prior: agedi.diffusion.distributions.Distribution = UniformCell(), sde: Optional[agedi.diffusion.sdes.SDE] = None, **kwargs) Bases: :py:obj:`agedi.diffusion.noisers.Noiser` Implements noising of atoms positions in Cartesian coordinates. :param sde_class: The class of the SDE to be used for the noising. :type sde_class: SDE :param sde_kwargs: The keyword arguments to be passed to the SDE class. :type sde_kwargs: Dict :param distribution: The distribution to be used for the noise. :type distribution: Distribution :param prior: The prior distribution to be used for the noise. :type prior: Distribution :param sde: An already-instantiated SDE object. When provided, *sde_class* and *sde_kwargs* are ignored. Useful for reconstructing a noiser from saved hyperparameters. :type sde: SDE, optional :param key: The key to be used for the noising. :type key: str :param \*\*kwargs: Additional keyword arguments to be passed to the Noiser class. :returns: The noiser for the atoms positions in Cartesian coordinates. :rtype: Noiser .. py:attribute:: _key :value: 'pos' .. py:method:: get_hparams() -> Dict Return hyperparameters for this positions noiser. .. py:method:: _noise(batch: agedi.data.AtomsGraph) -> agedi.data.AtomsGraph Initializes the noise for the positions noiser. Added noise is stored in the self.key+"_noise", which by default is "positions_noise". :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 positions of the atomistic structure. The denoising follows the Euler-Maruyama scheme. ::math:: R_i+1 = R_i + \Delta t (f(R_i, t) + g(t)**2 * s(R_i, t)) + \sqrt{\Delta t} g(t) * w The used score is expected to be stored in the self.key+"_score", which by default is "pos_score". :param batch: The atomistic structure (or batch hereof) to be denoised. :type batch: AtomsGraph :param delta_t: The time step 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) -> torch.Tensor Compute the noiser loss. Computes the loss of the diffusion model for the positions noiser Expects the total added positions noise to be stored in the self.key+"_noise", which by default is "pos_noise" and the predicted score to be stored in the self.key+"_score", which by default is "pos_score". The loss is computed as ::math:: L = \sum_i ||\sigma_t w_i + \sigma_t^2 s(R_i)||^2 With the noise taking into account periodic boundary conditions. :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:: periodic_distance(X: torch.tensor, N: torch.tensor, cells: torch.tensor, idxs: torch.tensor) -> torch.tensor Periodic distance computation. Takes X and N (noise) and computes the minimum distance between X and Y=X+N taking into account periodic boundary conditions. :param X: The positions (N, 3) :type X: torch.Tensor :param N: The noise (N, 3) :type N: torch.Tensor :param cell: The cell (3*K, 3) :type cell: torch.Tensor :param idxs: The indices of atoms in graphs (N,) :type idxs: torch.Tensor :returns: **dist** -- The distance between X and Y=X+N :rtype: torch.Tensor .. py:class:: Positions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Optional[Dict] = None, sde: Optional[agedi.diffusion.sdes.SDE] = None, distribution: agedi.diffusion.distributions.Distribution = Normal(), prior: agedi.diffusion.distributions.Distribution = StandardNormal(), **kwargs) Bases: :py:obj:`PositionsNoiser` Positions noiser with :class:`~agedi.diffusion.distributions.StandardNormal` prior and :class:`~agedi.diffusion.distributions.Normal` noise distribution. This is the base positions noiser suited for gas-phase clusters or systems where positions are not constrained to a periodic unit cell. The SDE can still be chosen freely via the *sde* parameter. Subclasses can override the ``distribution`` and ``prior`` while still delegating to this class through ``super()``. :param sde_class: Class of the SDE to use. Defaults to :class:`~agedi.diffusion.sdes.VE`. Ignored when *sde* is provided. :type sde_class: SDE, optional :param sde_kwargs: Keyword arguments forwarded to *sde_class*. Ignored when *sde* is provided. :type sde_kwargs: dict, optional :param sde: Pre-instantiated SDE object. When provided *sde_class* and *sde_kwargs* are ignored. :type sde: SDE, optional :param distribution: Noise distribution. Subclasses may supply a different default. :type distribution: Distribution, optional :param prior: Prior distribution. Subclasses may supply a different default. :type prior: Distribution, optional :param \*\*kwargs: Additional keyword arguments forwarded to :class:`~agedi.diffusion.noisers.PositionsNoiser`. .. py:method:: get_hparams() -> Dict Return hyperparameters for this positions noiser. Only includes :attr:`sde` and :attr:`loss_scaling`; the distribution and prior are fixed by the class and not needed for reconstruction. .. py:class:: CellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Optional[Dict] = None, sde: Optional[agedi.diffusion.sdes.SDE] = None, **kwargs) Bases: :py:obj:`Positions` Positions noiser with :class:`~agedi.diffusion.distributions.UniformCell` prior and :class:`~agedi.diffusion.distributions.Normal` noise distribution. Suited for periodic bulk or surface systems where atoms should be initialised uniformly within the unit cell. Inherits from :class:`Positions`; the SDE can still be chosen freely. :param sde_class: Class of the SDE to use. Defaults to :class:`~agedi.diffusion.sdes.VE`. Ignored when *sde* is provided. :type sde_class: SDE, optional :param sde_kwargs: Keyword arguments forwarded to *sde_class*. Ignored when *sde* is provided. :type sde_kwargs: dict, optional :param sde: Pre-instantiated SDE object. :type sde: SDE, optional :param \*\*kwargs: Additional keyword arguments forwarded to :class:`~agedi.diffusion.noisers.PositionsNoiser`. .. py:class:: ConfinedCellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Optional[Dict] = None, sde: Optional[agedi.diffusion.sdes.SDE] = None, **kwargs) Bases: :py:obj:`Positions` Positions noiser with :class:`~agedi.diffusion.distributions.UniformCellConfined` prior and :class:`~agedi.diffusion.distributions.TruncatedNormal` noise distribution. Suited for surface adsorption or porous-material systems where atoms are confined to a Z-range within the unit cell. Inherits from :class:`Positions`; the SDE can still be chosen freely. :param sde_class: Class of the SDE to use. Defaults to :class:`~agedi.diffusion.sdes.VE`. Ignored when *sde* is provided. :type sde_class: SDE, optional :param sde_kwargs: Keyword arguments forwarded to *sde_class*. Ignored when *sde* is provided. :type sde_kwargs: dict, optional :param sde: Pre-instantiated SDE object. :type sde: SDE, optional :param \*\*kwargs: Additional keyword arguments forwarded to :class:`~agedi.diffusion.noisers.PositionsNoiser`. .. py:class:: Types(prior=Constant(0), distribution=Categorical(), noise_schedule: NoiseSchedule = NoiseSchedule(0.01, 3.0), sampling_mask: Optional[torch.Tensor] = None, n_classes: int = 100, type_map: Optional[List[int]] = None, **kwargs) Bases: :py:obj:`agedi.diffusion.noisers.Noiser` Type noiser. Based on score entropy and discrete diffusion model. See https://arxiv.org/abs/2310.16834 for more details. Uses an absorbing state (index 0) as the first state in the transition matrix. .. py:attribute:: _key :value: 'x' .. py:attribute:: noise_schedule .. py:attribute:: sampling_mask :value: None .. py:attribute:: n_classes :value: 100 .. py:method:: get_hparams() -> Dict Return hyperparameters for this types noiser. .. py:method:: _noise(batch: agedi.data.AtomsGraph) -> agedi.data.AtomsGraph Noises the attribute of the atomistic structure. Performs the noising of the atomic types. :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. Denoisis the atomic types. :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 last denoising step is performed. :type last: bool :returns: The denoised atomistic structure (or bach hereof). :rtype: AtomsGraph .. py:method:: _loss(batch: agedi.data.AtomsGraph) -> torch.Tensor Computes the training loss. The score is with score entropy training as thus given as score=log(s) and then for sampling should be used as a concrete score i.e. exp(score)! :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:: sample_transition(x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor Sample the transition vector for the types This corresponds to noising the types in the discrete diffusion model :param x: The current types :type x: torch.Tensor :param sigma: The total noise :type sigma: torch.Tensor :returns: The noised types :rtype: torch.Tensor .. py:method:: score_entropy(score: torch.Tensor, sigma: torch.Tensor, x: torch.Tensor, x0: torch.Tensor) -> torch.Tensor Computes the score entropy loss :param score: The score :type score: torch.Tensor :param sigma: The total noise :type sigma: torch.Tensor :param x: The noised types :type x: torch.Tensor :param x0: The original types :type x0: torch.Tensor :returns: The score entropy loss :rtype: torch.Tensor .. py:method:: transp_rate(x: torch.Tensor) -> torch.Tensor Compute the i'th row of the rate transition matrix Q Can be used to compute the reverse rate :param x: The types :type x: torch.Tensor :returns: The i'th row of the rate transition matrix Q :rtype: torch.Tensor .. py:method:: reverse_rate(x: torch.Tensor, score: torch.Tensor) -> torch.Tensor Constructs the reverse rate. The reverse rate is given as the score * transp_rate :param x: The types :type x: torch.Tensor :param score: The score :type score: torch.Tensor :returns: The reverse rate :rtype: torch.Tensor .. py:method:: sample_rate(callable: Callable, x: torch.Tensor, rate: torch.Tensor) -> torch.Tensor Sample the rate Explain more... :param callable: Callable function defining the categorical distribution :type callable: Callable :param x: The types :type x: torch.Tensor :param rate: The rate :type rate: torch.Tensor :returns: The sampled rate :rtype: torch.Tensor .. py:method:: staggered_score(score: torch.Tensor, dsigma: torch.Tensor) -> torch.Tensor Computes the staggered score Computes p_{sigma - dsigma}(z) / p_{sigma}(x), which is approximated with e^{-{dsigma} E} score :param score: The score :type score: torch.Tensor :param dsigma: The rate noise :type dsigma: torch.Tensor :returns: The staggered score :rtype: torch.Tensor .. py:method:: transp_transition(x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor Compute the transition matrix for the types Explain more.. :param x: The types :type x: torch.Tensor :param sigma: The total noise :type sigma: torch.Tensor :returns: The transition matrix :rtype: torch.Tensor .. py:data:: TypesNoiser