schrodinger.application.matsci.flywheel.ldmol.diffusion.utils module

Statistical utility functions for Gaussian diffusion.

schrodinger.application.matsci.flywheel.ldmol.diffusion.utils.normal_kl(mean1, logvar1, mean2, logvar2)

Compute KL divergence between two Gaussians.

Shapes are automatically broadcasted, so batches can be compared to scalars.

Equation:

KL(N(mu1, sigma1^2) || N(mu2, sigma2^2))
  = 0.5 * (-1 + logvar2 - logvar1
           + exp(logvar1 - logvar2)
           + (mu1 - mu2)^2 * exp(-logvar2))

Parameters:

• mean1 – Mean of the first distribution

• logvar1 – Log-variance of the first distribution

• mean2 – Mean of the second distribution

• logvar2 – Log-variance of the second distribution

Return type:

torch.Tensor

Returns:

KL divergence, same shape as inputs

schrodinger.application.matsci.flywheel.ldmol.diffusion.utils.approx_standard_normal_cdf(x)

Fast approximation of the standard normal CDF.

Uses the tanh approximation of the error function:

Phi(x) ≈ 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))

Parameters:

x (torch.Tensor) – Input values

Return type:

torch.Tensor

Returns:

Approximate CDF values in [0, 1]

schrodinger.application.matsci.flywheel.ldmol.diffusion.utils.continuous_gaussian_log_likelihood(x, *, means, log_scales)

Log-likelihood of a continuous Gaussian distribution.

Equation:

log p(x) = log N(x; mu, sigma^2)
         = log N((x - mu) / sigma; 0, 1)

where sigma = exp(log_scales).

Parameters:

• x (torch.Tensor) – Target values

• means (torch.Tensor) – Gaussian means

• log_scales (torch.Tensor) – Gaussian log standard deviations

Return type:

torch.Tensor

Returns:

Log probabilities in nats, same shape as x

schrodinger.application.matsci.flywheel.ldmol.diffusion.utils.discretized_gaussian_log_likelihood(x, *, means, log_scales)

Log-likelihood of a Gaussian discretized to an image grid.

Assumes pixel values are uint8 rescaled to [-1, 1]; each bin has width 1/255. Let z = (x - mu) / sigma:

log p(x) = log(Phi(z + 1/255*inv_sigma) - Phi(z - 1/255*inv_sigma))

Boundary cases:

x < -0.999:  log Phi(z + 1/255*inv_sigma)         (left tail)
x >  0.999:  log(1 - Phi(z - 1/255*inv_sigma))    (right tail)

Parameters:

• x (torch.Tensor) – Target values in [-1, 1]

• means (torch.Tensor) – Gaussian means

• log_scales (torch.Tensor) – Gaussian log standard deviations

Return type:

torch.Tensor

Returns:

Log probabilities in nats, same shape as x