agedi.diffusion.distributions.normal¶

Attributes¶

_CONFINEMENT_CLAMP_EPS

Classes¶

`StandardNormal`	Standard Normal Distribution
`Normal`	Normal Distribution
`TruncatedNormal`	Truncated Normal Distribution
`ZeroComNormal`	Normal distribution whose noise increment has zero center of mass per graph.
`ZeroComStandardNormal`	Standard Normal distribution with zero center of mass per graph.

Functions¶

_zero_com(→ torch.Tensor)

Subtract the per-graph center of mass from x.

Module Contents¶

agedi.diffusion.distributions.normal._zero_com(x: torch.Tensor, batch_idx: torch.Tensor) → torch.Tensor¶

Subtract the per-graph center of mass from x.

Projects x onto the zero-COM subspace: each graph’s mean column is subtracted so that the COM of the returned tensor is exactly zero for every graph. This is the translational-invariance projection from the EDM paper (Hoogeboom et al., NeurIPS 2022, arXiv:2203.17003).

Parameters:

x (torch.Tensor, shape (n_atoms, d))
batch_idx (torch.Tensor, shape (n_atoms,)) – Graph membership index (batch.batch).

Return type:

torch.Tensor, shape (n_atoms, d)

agedi.diffusion.distributions.normal._CONFINEMENT_CLAMP_EPS = 0.0001¶

class agedi.diffusion.distributions.normal.StandardNormal(scale: float = 1.0, **kwargs)¶

Bases: agedi.diffusion.distributions.Distribution

Standard Normal Distribution

Parameters:: scale (float, optional) – Standard deviation used when sampling. Defaults to 1.0. Set this to the SDE’s sigma_max (or sqrt(var(t=1))) so that the prior matches the forward-process marginal at T=1, replacing the old 0.8 * N**(1/3) heuristic which was arbitrary and broke for non-compact or heterogeneously-sized systems.

scale = 1.0¶

get_hparams() → dict¶: Return hyperparameters for this distribution.

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

Prepare the distribution for sampling from batch.

Sets self.shape to (n_atoms, *trailing) where n_atoms is read from batch.n_atoms and the trailing dimensions come from the existing attribute. Using n_atoms rather than the attribute’s leading dimension avoids a shape-mismatch when called during graph initialisation (via initialize_graph()), where the attribute tensor may still be empty even though n_atoms has already been set.

Parameters:: batch (AtomsGraph) – Batch of atomistic data.

_sample(shape: torch.Size | None = None, **kwargs) → torch.Tensor¶

Sample from the standard normal distribution.

Parameters:: shape (torch.Size, optional) – Output shape. Defaults to self.shape set during _setup.
Returns:: Sampled tensor with std equal to self.scale.
Return type:: torch.Tensor

class agedi.diffusion.distributions.normal.Normal¶

Bases: agedi.diffusion.distributions.Distribution

Normal Distribution

_sample(mu: torch.Tensor, sigma: torch.Tensor, **kwargs) → torch.Tensor¶

Sample from the normal distribution

Parameters:

mu (torch.Tensor) – Mean of the distribution
sigma (torch.Tensor) – Standard deviation of the distribution

Returns:

Sampled tensor

Return type:

torch.Tensor

class agedi.diffusion.distributions.normal.TruncatedNormal(index: int = 2, **kwargs)¶

Bases: agedi.diffusion.distributions.Distribution

Truncated Normal Distribution

Parameters:: index (int) – The index of the property to truncate

index = 2¶

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

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

Setup the distribution

Prepare the distribution for sampling of the batch

Parameters:: batch (AtomsGraph) – Batch of data
Return type:: None

_sample(mu: torch.Tensor, sigma: torch.Tensor, **kwargs) → torch.Tensor¶

Sample from the truncated normal distribution

Parameters:

mu (torch.Tensor) – Mean of the distribution
sigma (torch.Tensor) – Standard deviation of the distribution

Returns:

Sampled tensor

Return type:

torch.Tensor

class agedi.diffusion.distributions.normal.ZeroComNormal¶

Bases: Normal

Normal distribution whose noise increment has zero center of mass per graph.

Drop-in replacement for Normal for use with the Positions noiser on gas-phase molecules and clusters. After sampling x = N(mu, sigma), the COM of the noise increment (x - mu) is subtracted per graph so that the diffusion process operates in the translationally-invariant subspace.

Reference: Hoogeboom et al., “Equivariant Diffusion for Molecule Generation in 3D”, NeurIPS 2022. arXiv:2203.17003

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

_sample(mu: torch.Tensor, sigma: torch.Tensor, **kwargs) → torch.Tensor¶

Sample from the normal distribution

Parameters:

mu (torch.Tensor) – Mean of the distribution
sigma (torch.Tensor) – Standard deviation of the distribution

Returns:

Sampled tensor

Return type:

torch.Tensor

class agedi.diffusion.distributions.normal.ZeroComStandardNormal(scale: float = 1.0, **kwargs)¶

Bases: StandardNormal

Standard Normal distribution with zero center of mass per graph.

Drop-in replacement for StandardNormal intended as the prior distribution for the Positions noiser. Sampled positions are centered at the origin for every graph.

Parameters:

scale (float, optional) – Standard deviation of the prior. Should be set to sqrt(sde.var(t=1)) (i.e. sigma_max for a VE-SDE) so that the prior matches the forward-process marginal at T=1. Defaults to 1.0; Positions sets this automatically from the SDE.
Reference (Hoogeboom et al., "Equivariant Diffusion for Molecule Generation)
3D" (in)
arXiv (NeurIPS 2022.)

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

Prepare the distribution for sampling from batch.

Parameters:: batch (AtomsGraph) – Batch of atomistic data.

_sample(shape: torch.Size | None = None, **kwargs) → torch.Tensor¶

Sample from the standard normal distribution.

Parameters:: shape (torch.Size, optional) – Output shape. Defaults to self.shape set during _setup.
Returns:: Sampled tensor with std equal to self.scale.
Return type:: torch.Tensor