schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion module¶

Gaussian diffusion process for LDMol molecular generation.

Implements DDPM forward process, DDIM reverse sampling, and training losses for DPO fine-tuning.

schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.mean_flat(tensor)¶: Average over all non-batch dimensions.

schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.extract_into_tensor(arr, timesteps, broadcast_shape)¶

Index a 1-D numpy array by a batch of timestep indices and broadcast the result to broadcast_shape.

Parameters:

arr (numpy.ndarray) – 1-D array of precomputed schedule values
timesteps (torch.Tensor) – 1-D integer tensor of indices
broadcast_shape (tuple) – target shape (batch, …)

Return type:

torch.Tensor

Returns:

Values at timesteps, broadcast to broadcast_shape

class schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.ModelMeanType¶

Bases: Enum

What the model output represents.

PREVIOUS_X = 1¶

START_X = 2¶

EPSILON = 3¶

class schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.ModelVarType¶

Bases: Enum

How the model’s output variance is determined.

LEARNED = 1¶

FIXED_SMALL = 2¶

FIXED_LARGE = 3¶

LEARNED_RANGE = 4¶

schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.get_named_beta_schedule(num_diffusion_timesteps)¶

Return the linear beta schedule for the LDMol checkpoint.

Parameters:: num_diffusion_timesteps (int) – Number of diffusion steps (T)
Return type:: numpy.ndarray
Returns:: Beta values of shape (T,)

class schrodinger.application.matsci.flywheel.ldmol.diffusion.gaussian_diffusion.GaussianDiffusion(*, betas, model_mean_type, model_var_type)¶

Bases: object

DDPM forward process and DDIM reverse sampling.

The forward process adds noise according to a beta schedule:

q(x_t | x_{t-1}) = N(x_t; sqrt(1-beta_t)*x_{t-1}, beta_t*I)

Inference uses DDIM sampling to generate latents from noise conditioned on text embeddings via the DiT model.

__init__(*, betas, model_mean_type, model_var_type)¶

Parameters:

betas (numpy.ndarray) – Beta schedule of shape (T,)
model_mean_type (ModelMeanType) – What the model predicts
model_var_type (ModelVarType) – How variance is determined

weightedSum(coef_a, coef_b, tensor_a, tensor_b, timesteps)¶

Compute a schedule-weighted linear combination of two tensors.

Returns coef_a(t) * tensor_a + coef_b(t) * tensor_b, where the schedule coefficients are broadcast to tensor_a.shape.

Parameters:

coef_a (numpy.ndarray) – 1-D schedule array for tensor_a
coef_b (numpy.ndarray) – 1-D schedule array for tensor_b
tensor_a (torch.Tensor) – First tensor
tensor_b (torch.Tensor) – Second tensor
timesteps (torch.Tensor) – Timestep indices (batch,)

Return type:

torch.Tensor

Returns:

Weighted sum, same shape as tensor_a

qMeanVariance(x_start, timesteps)¶

Compute the mean and variance of q(x_t | x_0).

Parameters:

x_start (torch.Tensor) – Clean latent x_0, shape (batch, …)
timesteps (torch.Tensor) – Timestep indices, shape (batch,)

Return type:

tuple(torch.Tensor, torch.Tensor)

Returns:

mean, log_variance each of shape (batch, …)

qSample(x_start, timesteps, noise=None)¶

Sample from q(x_t | x_0) — add noise to x_start for t steps.

Parameters:

x_start (torch.Tensor) – Clean latent x_0
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
noise (torch.Tensor) – Gaussian noise; sampled if None

Return type:

torch.Tensor

Returns:

Noisy latent x_t of same shape as x_start

qPosteriorMeanVariance(x_start, noisy, timesteps)¶

Compute mean and variance of the posterior q(x_{t-1} | x_t, x_0).

Parameters:

x_start (torch.Tensor) – Clean latent x_0
noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)

Return type:

tuple(torch.Tensor, torch.Tensor)

Returns:

posterior_mean, posterior_log_variance_clipped

predictXstartFromEps(noisy, timesteps, eps)¶

Recover x_0 from the predicted noise eps.

Parameters:

noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
eps (torch.Tensor) – Predicted noise

Return type:

torch.Tensor

Returns:

Predicted clean latent x_0

predictEpsFromXstart(noisy, timesteps, pred_xstart)¶

Recover the predicted noise from x_0 and x_t.

Parameters:

noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
pred_xstart (torch.Tensor) – Predicted clean latent x_0

Return type:

torch.Tensor

Returns:

Implied noise prediction

pMeanVariance(model, noisy, timesteps, clip_denoised=True, denoised_fn=None, model_kwargs=None)¶

Apply the model to get p(x_{t-1} | x_t) mean and variance.

Assumes LEARNED_RANGE variance and EPSILON mean type (LDMol checkpoint).

Parameters:

model – DiT model; called with noisy, timesteps, and model_kwargs
noisy (torch.Tensor) – Noisy latent x_t, shape (batch, channels, …)
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
clip_denoised (bool) – Clamp predicted x_0 to [-1, 1]
denoised_fn – Optional function applied to pred_xstart before clipping
model_kwargs (dict) – Extra keyword arguments passed to model

Return type:

dict

Returns:

Keys ‘mean’, ‘variance’, ‘log_variance’, ‘pred_xstart’

conditionMean(cond_fn, p_mean_var, noisy, timesteps, model_kwargs=None)¶

Apply classifier guidance to the mean using Sohl-Dickstein et al. (2015).

Parameters:

cond_fn – Gradient function; returns grad(log p(y|x_t))
p_mean_var (dict) – Output of pMeanVariance
noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
model_kwargs (dict) – Extra kwargs forwarded to cond_fn

Return type:

torch.Tensor

Returns:

Conditioned mean

conditionScore(cond_fn, p_mean_var, noisy, timesteps, model_kwargs=None)¶

Apply classifier guidance to the score using Song et al. (2020).

Parameters:

cond_fn – Gradient function; returns grad(log p(y|x_t))
p_mean_var (dict) – Output of pMeanVariance
noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
model_kwargs (dict) – Extra kwargs forwarded to cond_fn

Return type:

dict

Returns:

Updated p_mean_var with conditioned mean and pred_xstart

ddimSample(model, noisy, timesteps, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, eta=0.0)¶

Sample x_{t-1} from the model using one DDIM step.

Parameters:

model – DiT model
noisy (torch.Tensor) – Current noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
clip_denoised (bool) – Clamp predicted x_0 to [-1, 1]
denoised_fn – Optional function applied to pred_xstart
cond_fn – Optional classifier guidance gradient function
model_kwargs (dict) – Extra kwargs forwarded to model
eta (float) – Stochasticity scale; 0.0 = deterministic DDIM

Return type:

dict

Returns:

Keys ‘sample’ (x_{t-1}) and ‘pred_xstart’ (predicted x_0)

ddimSampleLoopProgressive(model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, eta=0.0)¶

Run DDIM reverse sampling, yielding each intermediate step.

Parameters:

model – DiT model
shape (tuple) – Output shape (batch, channels, seq_len)
noise (torch.Tensor) – Initial noise; sampled if None
clip_denoised (bool) – Clamp predicted x_0 to [-1, 1]
denoised_fn – Optional function applied to pred_xstart
cond_fn – Optional classifier guidance gradient function
model_kwargs (dict) – Extra kwargs forwarded to model
device – Compute device; inferred from model if None
eta (float) – Stochasticity scale; 0.0 = deterministic DDIM

Return type:

generator

Returns:

Yields dict with ‘sample’ and ‘pred_xstart’ at each step

ddimSampleLoop(model, shape, noise=None, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None, eta=0.0)¶

Generate samples from the model using DDIM.

Parameters:

model – DiT model
shape (tuple) – Output shape (batch, channels, seq_len)
noise (torch.Tensor) – Initial noise; sampled if None
clip_denoised (bool) – Clamp predicted x_0 to [-1, 1]
denoised_fn – Optional function applied to pred_xstart
cond_fn – Optional classifier guidance gradient function
model_kwargs (dict) – Extra kwargs forwarded to model
device – Compute device; inferred from model if None
eta (float) – Stochasticity scale; 0.0 = deterministic DDIM

Return type:

torch.Tensor

Returns:

Denoised latent of shape

vbTermsBpd(model, x_start, noisy, timesteps, clip_denoised=True, model_kwargs=None)¶

Compute one term of the variational lower bound in bits per dim.

Returns the KL divergence at t>0 and the decoder NLL at t=0.

Parameters:

model – DiT model (or a frozen lambda for variance-only training)
x_start (torch.Tensor) – Clean latent x_0
noisy (torch.Tensor) – Noisy latent x_t
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
clip_denoised (bool) – Clamp predicted x_0 to [-1, 1]
model_kwargs (dict) – Extra kwargs forwarded to model

Return type:

dict

Returns:

Keys ‘output’ (shape batch,) and ‘pred_xstart’

trainingLosses(model, x_start, timesteps, model_kwargs=None, noise=None)¶

Compute MSE + variational bound training losses for one timestep.

The VB term trains the variance head (LEARNED_RANGE) without affecting the epsilon gradient; model_output is detached before passing to vbTermsBpd.

Parameters:

model – DiT model
x_start (torch.Tensor) – Clean latent x_0
timesteps (torch.Tensor) – Timestep indices, shape (batch,)
model_kwargs (dict) – Extra kwargs forwarded to model
noise (torch.Tensor) – Gaussian noise; sampled if None

Return type:

dict

Returns:

Keys ‘mse’, ‘vb’, ‘loss’ each of shape (batch,)

scaleTimesteps(timesteps)¶

Return timesteps unchanged; SpacedDiffusion overrides this.

Parameters:: timesteps (torch.Tensor) – Timestep indices
Return type:: torch.Tensor
Returns:: Unchanged timesteps