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 `ZeroComStandardNormal`
`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, 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.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.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.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.

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¶