MolecularDiffusion.modules.models.diffpharma.conditional_ddpm¶

Pocket-conditioned DDPM for DiffPharma.

Port of ConditionalDDPM from others/DiffPharma/equivariant_diffusion/conditional_model.py (EDM / DiffSBDD lineage): epsilon parametrisation, predefined polynomial noise schedule, only the ligand is noised while pocket and both pharmacophore node sets stay clean and act as fixed context.

Dropped from the upstream file (all dead or broken in this repo’s config): SimpleConditionalDDPM, GammaNetwork (learned schedule), partially_noised_ligand / sample_p_zt_given_zs / sample_combined_position_feature_noise / sample_center_gravity_zero_gaussian_batch – all of which call pre-DiffPharma 2-node-set signatures and would raise on entry.

One upstream bug is fixed: the CoM-drift correction at the end of sample_given_pocket called remove_mean_batch with 4 of its 8 arguments, i.e. it crashed whenever it triggered.

Classes¶

`ConditionalDDPM`	Conditional diffusion module (pocket + pharmacophore particles as context).
`PredefinedNoiseSchedule`	Lookup table for a non-learned noise schedule.

Functions¶

`clip_noise_schedule`(alphas2[, clip_value])
`cosine_beta_schedule`(timesteps[, s, raise_to_power])
`polynomial_schedule`(timesteps[, s, power])	Noise schedule `1 - x^power`.

Module Contents¶

class MolecularDiffusion.modules.models.diffpharma.conditional_ddpm.ConditionalDDPM(dynamics, atom_nf, residue_nf, interh_nf, interhp_nf, n_dims, size_histogram, timesteps=1000, parametrization='eps', noise_schedule='polynomial_2', noise_precision=0.0001, loss_type='l2', norm_values=(1.0, 1.0), norm_biases=(None, 0.0), virtual_node_idx=None)¶

Bases: torch.nn.Module

Conditional diffusion module (pocket + pharmacophore particles as context).

static SNR(gamma)¶

alpha(gamma, target_tensor)¶

static assert_mean_zero_with_mask(x, node_mask, eps=1e-10)¶

static cdf_standard_gaussian(x)¶

check_issues_norm_values(num_stdevs=8)¶

compute_x_pred(net_out, zt, gamma_t, batch_mask)¶

delta_log_px(num_nodes)¶

forward(ligand, pocket, interh, interhp, return_info=False)¶: Compute the loss terms (see the task module for how they combine).

static gaussian_KL(q_mu_minus_p_mu_squared, q_sigma, p_sigma, d)¶

static inflate_batch_array(array, target)¶

kl_prior(xh_lig, mask_lig, num_nodes)¶

log_constants_p_x_given_z0(n_nodes, device)¶

log_pN(N_lig, N_pocket)¶

log_pxh_given_z0_without_constants(ligand, z_0_lig, eps_lig, net_out_lig, gamma_0, epsilon=1e-10)¶

noised_representation(xh_lig, xh0_pocket, xh0_interh, xh0_interhp, lig_mask, pocket_mask, interh_mask, interhp_mask, gamma_t)¶

normalize(ligand=None, pocket=None)¶

classmethod remove_mean_batch(x_lig, x_pocket, x_interh, x_interhp, lig_indices, pocket_indices, interh_indices, interhp_indices)¶: Subtract the ligand centre of mass from every node set.

abstractmethod sample(*args, **kwargs)¶

static sample_gaussian(size, device)¶

sample_given_pocket(pocket, interh, interhp, num_nodes_lig, return_frames=1, timesteps=None)¶

Reverse-diffuse a ligand inside pocket.

Returns (out_lig, out_pocket, lig_mask, pocket_mask) – flat, scatter-masked tensors, NOT (B, N, .).

sample_normal_zero_com(mu_lig, xh0_pocket, xh0_interh, xh0_interhp, sigma, lig_mask, pocket_mask, interh_mask, interhp_mask, fix_noise=False)¶

sample_p_xh_given_z0(z0_lig, xh0_pocket, xh0_interh, xh0_interhp, lig_mask, pocket_mask, interh_mask, interhp_mask, batch_size, fix_noise=False)¶

sample_p_zs_given_zt(s, t, zt_lig, xh0_pocket, xh0_interh, xh0_interhp, ligand_mask, pocket_mask, interh_mask, interhp_mask, fix_noise=False)¶

sigma(gamma, target_tensor)¶

sigma_and_alpha_t_given_s(gamma_t, gamma_s, target_tensor)¶

subspace_dimensionality(input_size)¶

static sum_except_batch(x, indices)¶

unnormalize(x, h_cat)¶

unnormalize_z(z_lig, z_pocket)¶

xh_given_zt_and_epsilon(z_t, epsilon, gamma_t, batch_mask)¶: Equation (7) in the EDM paper.

T = 1000¶

atom_nf¶

dynamics¶

gamma¶

interh_nf¶

interhp_nf¶

loss_type = 'l2'¶

n_dims¶

norm_biases = (None, 0.0)¶

norm_values = (1.0, 1.0)¶

parametrization = 'eps'¶

residue_nf¶

size_distribution¶

vnode_idx = None¶

class MolecularDiffusion.modules.models.diffpharma.conditional_ddpm.PredefinedNoiseSchedule(noise_schedule, timesteps, precision)¶

Bases: torch.nn.Module

Lookup table for a non-learned noise schedule.

forward(t)¶

gamma¶

timesteps¶

MolecularDiffusion.modules.models.diffpharma.conditional_ddpm.clip_noise_schedule(alphas2, clip_value=0.001)¶

MolecularDiffusion.modules.models.diffpharma.conditional_ddpm.cosine_beta_schedule(timesteps, s=0.008, raise_to_power: float = 1)¶

MolecularDiffusion.modules.models.diffpharma.conditional_ddpm.polynomial_schedule(timesteps: int, s=0.0001, power=3.0)¶: Noise schedule 1 - x^power.