agedi.diffusion.noisers.pos =========================== .. py:module:: agedi.diffusion.noisers.pos Classes ------- .. autoapisummary:: agedi.diffusion.noisers.pos.PositionsNoiser agedi.diffusion.noisers.pos.Positions agedi.diffusion.noisers.pos.CellPositions agedi.diffusion.noisers.pos.ConfinedCellPositions Module Contents --------------- .. 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`.