MolecularDiffusion.modules.models.difflinker.edm

E(n)-equivariant fragment/linker diffusion model, ported near-verbatim from DiffLinker’s src/edm.py.

Only EDM’s non-inpainting path is exercised in this integration — InpaintingEDM ports as-is but stays unexercised (out of scope, see docs/model_integrations/difflinker/INTEGRATION_PLAN.md).

Classes

EDM
InpaintingEDM	Soft/probabilistic inpainting variant — ports as-is but stays

Module Contents

class MolecularDiffusion.modules.models.difflinker.edm.EDM(dynamics: MolecularDiffusion.modules.models.difflinker.egnn.Dynamics, in_node_nf: int, n_dims: int, timesteps: int = 1000, noise_schedule='learned', noise_precision=0.0001, loss_type='vlb', norm_values=(1.0, 1.0, 1.0), norm_biases=(None, 0.0, 0.0))

Bases: torch.nn.Module

SNR(gamma)

alpha(gamma, target_tensor)

static cdf_standard_gaussian(x)

compute_x_pred(eps_t, z_t, gamma_t)

Computes the most likely prediction of x.

delta_log_px(mask)

dimensionality(mask)

static expm1(x: torch.Tensor) → torch.Tensor

forward(x, h, node_mask, fragment_mask, linker_mask, edge_mask, context=None)

static gaussian_kl(q_mu, q_sigma, p_mu, p_sigma)

static gaussian_kl_for_dimension(q_mu, q_sigma, p_mu, p_sigma, d)

static inflate_batch_array(array, target)

kl_prior(xh, mask)

KL between q(z1 | x) and the prior p(z1) = Normal(0, 1).

log_constant_of_p_x_given_z0(x, mask)

log_p_xh_given_z0_without_constants(h, z_0, gamma_0, eps, eps_hat, mask, epsilon=1e-10)

normalize(x, h)

static numbers_of_nodes(mask)

sample_chain(x, h, node_mask, fragment_mask, linker_mask, edge_mask, context, keep_frames=None)

sample_combined_position_feature_noise(n_samples, n_nodes, mask)

sample_normal(mu, sigma, node_mask)

sample_p_xh_given_z0_only_linker(z_0, node_mask, fragment_mask, linker_mask, edge_mask, context)

Samples x ~ p(x|z0). Samples only linker features and coords.

sample_p_zs_given_zt_only_linker(s, t, z_t, node_mask, fragment_mask, linker_mask, edge_mask, context)

Samples zs ~ p(zs | zt). Only used during sampling. Samples only linker features and coords.

sigma(gamma, target_tensor)

sigma_and_alpha_t_given_s(gamma_t: torch.Tensor, gamma_s: torch.Tensor, target_tensor: torch.Tensor)

alpha t given s = alpha t / alpha s sigma t given s = sqrt(1 - (alpha t given s)^2)

static softplus(x: torch.Tensor) → torch.Tensor

static sum_except_batch(x)

unnormalize(x, h)

unnormalize_z(z)

T = 1000

dynamics

in_node_nf

n_dims

norm_biases = (None, 0.0, 0.0)

norm_values = (1.0, 1.0, 1.0)

class MolecularDiffusion.modules.models.difflinker.edm.InpaintingEDM(dynamics: MolecularDiffusion.modules.models.difflinker.egnn.Dynamics, in_node_nf: int, n_dims: int, timesteps: int = 1000, noise_schedule='learned', noise_precision=0.0001, loss_type='vlb', norm_values=(1.0, 1.0, 1.0), norm_biases=(None, 0.0, 0.0))

Bases: EDM

Soft/probabilistic inpainting variant — ports as-is but stays unexercised this pass (out of scope, see INTEGRATION_PLAN.md).

dimensionality(mask)

forward(x, h, node_mask, fragment_mask, linker_mask, edge_mask, context=None)

sample_chain(x, h, node_mask, edge_mask, fragment_mask, linker_mask, context, keep_frames=None)

sample_combined_position_feature_noise(n_samples, n_nodes, mask)

sample_p_xh_given_z0(z_0, node_mask, edge_mask, context)

sample_p_zs_given_zt(s, t, z_t, node_mask, edge_mask, context)

sample_q_xh_given_z0_and_x(z_0, node_mask)

sample_q_zs_given_zt_and_x(s, t, z_t, x, node_mask)