agedi.diffusion.noisers.pos

Classes

PositionsNoiser

Implements noising of atoms positions in Cartesian coordinates.

Positions

Positions noiser with ZeroComStandardNormal

CellPositions

Positions noiser with UniformCell prior

ConfinedCellPositions

Positions noiser with UniformCellConfined

Module Contents

class agedi.diffusion.noisers.pos.PositionsNoiser(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Dict | None = None, distribution: agedi.diffusion.distributions.Distribution = Normal(), prior: agedi.diffusion.distributions.Distribution = UniformCell(), sde: agedi.diffusion.sdes.SDE | None = None, loss_weighting: str = 'uniform', prediction_type: str = 'score', denoising_step: str = 'em', sampler: str | None = None, **kwargs)

Bases: agedi.diffusion.noisers.Noiser

Implements noising of atoms positions in Cartesian coordinates.

Parameters:

  • sde_class (SDE) – The class of the SDE to be used for the noising.

  • sde_kwargs (Dict) – The keyword arguments to be passed to the SDE class.

  • distribution (Distribution) – The distribution to be used for the noise.

  • prior (Distribution) – The prior distribution to be used for the noise.

  • sde (SDE, optional) – An already-instantiated SDE object. When provided, sde_class and sde_kwargs are ignored. Useful for reconstructing a noiser from saved hyperparameters.

  • key (str) – The key to be used for the noising.

  • **kwargs – Additional keyword arguments to be passed to the Noiser class.

Returns:

The noiser for the atoms positions in Cartesian coordinates.

Return type:

Noiser

_key = 'pos'
loss_weighting = 'uniform'
prediction_type = 'score'
sampler = 'em'
get_hparams() → Dict

Return hyperparameters for this positions noiser.

_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”.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be noised.

Returns:

The noised atomistic structure (or bach hereof).

Return type:

AtomsGraph

_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”.

Parameters:

  • batch (AtomsGraph) – The atomistic structure (or batch hereof) to be denoised.

  • delta_t (float) – The time step for the denoising.

  • last (bool) – If the denoising is the last step of the denoising.

Returns:

The denoised atomistic structure (or bach hereof).

Return type:

AtomsGraph

_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.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be noised and denoised.

Returns:

The loss of the noised and denoised atomistic structure.

Return type:

float

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.

Parameters:

  • X (torch.Tensor) – The positions (N, 3)

  • N (torch.Tensor) – The noise (N, 3)

  • cell (torch.Tensor) – The cell (3*K, 3)

  • idxs (torch.Tensor) – The indices of atoms in graphs (N,)

Returns:

dist – The distance between X and Y=X+N

Return type:

torch.Tensor

class agedi.diffusion.noisers.pos.Positions(sde_class: agedi.diffusion.sdes.SDE = VP, sde_kwargs: Dict | None = None, sde: agedi.diffusion.sdes.SDE | None = None, distribution: agedi.diffusion.distributions.Distribution = ZeroComNormal(), prior: agedi.diffusion.distributions.Distribution | None = None, **kwargs)

Bases: PositionsNoiser

Positions noiser with ZeroComStandardNormal prior and ZeroComNormal 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().

When prior is not supplied, the prior scale is set automatically to sqrt(sde.var(t=1)) — equal to sigma_max for a VE-SDE — so that the prior matches the forward-process marginal at T=1.

Parameters:

  • sde_class (SDE, optional) – Class of the SDE to use. Defaults to VP. Ignored when sde is provided.

  • sde_kwargs (dict, optional) – Keyword arguments forwarded to sde_class. Ignored when sde is provided.

  • sde (SDE, optional) – Pre-instantiated SDE object. When provided sde_class and sde_kwargs are ignored.

  • distribution (Distribution, optional) – Noise distribution. Subclasses may supply a different default.

  • prior (Distribution, optional) – Prior distribution. When None (default), a ZeroComStandardNormal with scale = sqrt(sde.var(t=1)) is created automatically.

  • **kwargs – Additional keyword arguments forwarded to PositionsNoiser.

class agedi.diffusion.noisers.pos.CellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Dict | None = None, sde: agedi.diffusion.sdes.SDE | None = None, distribution: agedi.diffusion.distributions.Distribution = Normal(), prior: agedi.diffusion.distributions.Distribution = UniformCell(), **kwargs)

Bases: Positions

Positions noiser with UniformCell prior and Normal noise distribution.

Suited for periodic bulk or surface systems where atoms should be initialised uniformly within the unit cell. Inherits from Positions; the SDE can still be chosen freely.

Parameters:

  • sde_class (SDE, optional) – Class of the SDE to use. Defaults to VE. Ignored when sde is provided.

  • sde_kwargs (dict, optional) – Keyword arguments forwarded to sde_class. Ignored when sde is provided.

  • sde (SDE, optional) – Pre-instantiated SDE object.

  • **kwargs – Additional keyword arguments forwarded to PositionsNoiser.

get_hparams() → Dict

Return hyperparameters for this noiser.

Distribution and prior are class-fixed defaults and excluded so that Hydra round-trip instantiation does not conflict with the explicit constructor defaults.

class agedi.diffusion.noisers.pos.ConfinedCellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Dict | None = None, sde: agedi.diffusion.sdes.SDE | None = None, distribution: agedi.diffusion.distributions.Distribution = TruncatedNormal(), prior: agedi.diffusion.distributions.Distribution = UniformCellConfined(), **kwargs)

Bases: Positions

Positions noiser with UniformCellConfined prior and 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 Positions; the SDE can still be chosen freely.

Parameters:

  • sde_class (SDE, optional) – Class of the SDE to use. Defaults to VE. Ignored when sde is provided.

  • sde_kwargs (dict, optional) – Keyword arguments forwarded to sde_class. Ignored when sde is provided.

  • sde (SDE, optional) – Pre-instantiated SDE object.

  • **kwargs – Additional keyword arguments forwarded to PositionsNoiser.

get_hparams() → Dict

Return hyperparameters for this noiser.

Distribution and prior are class-fixed defaults and excluded so that Hydra round-trip instantiation does not conflict with the explicit constructor defaults.