agedi.diffusion.noisers¶

Submodules¶

Attributes¶

TypesNoiser

Classes¶

`Noiser`	Noiser Base class
`SDENoiser`	Implements a SDE base class that can be inherited by other classes.
`PositionsNoiser`	Implements noising of atoms positions in Cartesian coordinates.
`Positions`	Positions noiser with `StandardNormal` prior
`CellPositions`	Positions noiser with `UniformCell` prior
`ConfinedCellPositions`	Positions noiser with `UniformCellConfined`
`Types`	Type noiser.

Package Contents¶

class agedi.diffusion.noisers.Noiser(distribution: agedi.diffusion.distributions.Distribution, prior: agedi.diffusion.distributions.Distribution, loss_scaling: float = 1.0, key: str | None = None, **kwargs)¶

Bases: abc.ABC, torch.nn.Module

Noiser Base class

Impments a noiser that can noise and denoise a atomistic structure attribute.

Parameters:

distribution (Distribution) – The distribution to be used for the noising.
prior (Distribution) – The prior to be used for the denoising.
loss_scaling (float) – Scaling factor applied to this noiser’s loss contribution.
key (str, optional) – 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.

Return type:

Noiser

_key: str¶

_registry: ClassVar[Dict[str, Callable[..., Noiser]]]¶

distribution¶

prior¶

loss_scaling = 1.0¶

classmethod register(name: str, factory: Callable[..., Noiser]) → None¶

The factory is called with sde as a keyword argument containing the resolved SDE instance. Noisers that do not use an SDE can safely ignore it via **kwargs.

Parameters:

name (str) – Alias string used to reference the noiser in create_diffusion().
factory (Callable) – A callable that accepts sde as a keyword argument and returns a Noiser instance.

Examples

from agedi.diffusion.noisers import Noiser

class MyNoiser(Noiser):
    ...

Noiser.register("my_noiser", lambda sde: MyNoiser(sde=sde))

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.
Return type:: dict

property key: str¶: The key of the attribute to be noised and denoised.

abstractmethod _noise(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Noises the attribute of the atomistic structure.

Must be implemented by the subclass.

Parameters:: batch (AtomsGraph) – The atomistic structure (or batch hereof) to be noised.
Returns:: The noised atomistic structure (or bach hereof).
Return type:: AtomsGraph

abstractmethod _denoise(batch: agedi.data.AtomsGraph, delta_t: float, last: bool) → agedi.data.AtomsGraph¶

Denoises the attribute of the atomistic structure.

Must be implemented by the subclass.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be denoised.
delta_t (float) – The time step to be used for the denoising.

Returns:

The denoised atomistic structure (or bach hereof).

Return type:

AtomsGraph

abstractmethod _loss(batch: agedi.data.AtomsGraph) → float¶

Computes the training loss.

Must be implemented by the subclass.

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

noise(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Noises the attribute of the atomistic structure.

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 attribute of the atomistic structure.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be denoised.
delta_t (float) – The time step to be used 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) → float¶

Compute the training loss.

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

langevin_step(batch: agedi.data.AtomsGraph, step_size: 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 _denoise() with last=False and a fixed step_size. Subclasses may override this for a more specialised corrector.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be corrected.
step_size (float or torch.Tensor, optional) – Size of the Langevin corrector step. Passing a pre-created torch.Tensor avoids repeated tensor allocation when this method is called in a tight loop. Defaults to 0.01.

Returns:

The corrected atomistic structure.

Return type:

AtomsGraph

initialize_graph(batch: agedi.data.AtomsGraph) → None¶

Initializes the graph with the prior distribution.

Can be overwritten by subclasses for specific initializations.

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

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

Bases: agedi.diffusion.noisers.Noiser, abc.ABC

Implements a SDE base class that can be inherited by other classes.

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.
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 = None¶

get_hparams() → Dict¶: Return hyperparameters for this SDE noiser.

abstractmethod postprocess_score(score: torch.Tensor) → torch.Tensor¶

Post-process the predicted score before computing the loss.

Parameters:: score (torch.Tensor) – Raw predicted score tensor.
Returns:: Post-processed score tensor.
Return type:: torch.Tensor

abstractmethod postprocess_noise(noise: torch.Tensor) → torch.Tensor¶

Post-process the noise tensor before computing the loss.

Parameters:: noise (torch.Tensor) – Raw noise tensor.
Returns:: Post-processed noise tensor.
Return type:: torch.Tensor

_noise(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Adds noise to the atomistic structure.

Added noise is stored in the self.key+”_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”.

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

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

class agedi.diffusion.noisers.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, **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'¶

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.Positions(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 = StandardNormal(), **kwargs)¶

Bases: PositionsNoiser

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

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. 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. Subclasses may supply a different default.
**kwargs – Additional keyword arguments forwarded to PositionsNoiser.

get_hparams() → Dict¶

Return hyperparameters for this positions noiser.

Only includes sde and loss_scaling; the distribution and prior are fixed by the class and not needed for reconstruction.

class agedi.diffusion.noisers.CellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Dict | None = None, sde: agedi.diffusion.sdes.SDE | None = None, **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.

class agedi.diffusion.noisers.ConfinedCellPositions(sde_class: agedi.diffusion.sdes.SDE = VE, sde_kwargs: Dict | None = None, sde: agedi.diffusion.sdes.SDE | None = None, **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.

class agedi.diffusion.noisers.Types(prior=Constant(0), distribution=Categorical(), noise_schedule: NoiseSchedule = NoiseSchedule(0.01, 3.0), sampling_mask: torch.Tensor | None = None, n_classes: int = 100, type_map: List[int] | None = None, **kwargs)¶

Bases: 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.

_key = 'x'¶

noise_schedule¶

sampling_mask = None¶

n_classes = 100¶

get_hparams() → Dict¶: Return hyperparameters for this types noiser.

_noise(batch: agedi.data.AtomsGraph) → agedi.data.AtomsGraph¶

Noises the attribute of the atomistic structure.

Performs the noising of the atomic types.

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 attribute of the atomistic structure.

Denoisis the atomic types.

Parameters:

batch (AtomsGraph) – The atomistic structure (or batch hereof) to be denoised.
delta_t (float) – The time step to be used for the denoising.
last (bool) – If the last denoising step is performed.

Returns:

The denoised atomistic structure (or bach hereof).

Return type:

AtomsGraph

_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)!

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

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

Parameters:

x (torch.Tensor) – The current types
sigma (torch.Tensor) – The total noise

Returns:

The noised types

Return type:

torch.Tensor

score_entropy(score: torch.Tensor, sigma: torch.Tensor, x: torch.Tensor, x0: torch.Tensor) → torch.Tensor¶

Computes the score entropy loss

Parameters:

score (torch.Tensor) – The score
sigma (torch.Tensor) – The total noise
x (torch.Tensor) – The noised types
x0 (torch.Tensor) – The original types

Returns:

The score entropy loss

Return type:

torch.Tensor

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

Parameters:: x (torch.Tensor) – The types
Returns:: The i’th row of the rate transition matrix Q
Return type:: torch.Tensor

reverse_rate(x: torch.Tensor, score: torch.Tensor) → torch.Tensor¶

Constructs the reverse rate.

The reverse rate is given as the score * transp_rate

Parameters:

x (torch.Tensor) – The types
score (torch.Tensor) – The score

Returns:

The reverse rate

Return type:

torch.Tensor

sample_rate(callable: Callable, x: torch.Tensor, rate: torch.Tensor) → torch.Tensor¶

Sample the rate

Explain more…

Parameters:

callable (Callable) – Callable function defining the categorical distribution
x (torch.Tensor) – The types
rate (torch.Tensor) – The rate

Returns:

The sampled rate

Return type:

torch.Tensor

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

Parameters:

score (torch.Tensor) – The score
dsigma (torch.Tensor) – The rate noise

Returns:

The staggered score

Return type:

torch.Tensor

transp_transition(x: torch.Tensor, sigma: torch.Tensor) → torch.Tensor¶

Compute the transition matrix for the types

Explain more..

Parameters:

x (torch.Tensor) – The types
sigma (torch.Tensor) – The total noise

Returns:

The transition matrix

Return type:

torch.Tensor

agedi.diffusion.noisers.TypesNoiser¶