MolecularDiffusion.modules.models.diffsmol.score_model¶

DiffSMol hybrid diffusion score model, bond-free.

Ported from DiffSMol source/models/molopt_score_model.py (+ source/models/diffusion.py for the beta schedules).

Hybrid objective, both kept:

continuous Gaussian DDPM on positions (sigmoid beta schedule),
discrete D3PM categorical diffusion on atom types (cosine schedule, uniform or absorbing-mask transition kernel).

Dropped per the approved INTEGRATION_PLAN: the bond-type CE loss, the bond-distance loss, the bond-angle loss and the torsion-angle loss (the last two were already switched off by upstream’s own shipped config). The loss therefore reduces to loss_pos + loss_v_weight * loss_v.

Kept: shape conditioning and classifier-free guidance. Training zeroes the whole (B, 128, 3) latent with probability cond_mask_prob; that is what creates the unconditional branch, and it is why sample() can fall back to a zero latent when no reference shape is supplied.

Frame convention: center_pos_mode is 'none'. Coordinates are assumed to already be in the surface-centered frame (the caller subtracts the precomputed shape_center), which is the frame the shape latent lives in. Do not re-center here.

Attributes¶

DEFAULT_CONFIG

Classes¶

ScorePosNet3D

Shape-conditioned hybrid (Gaussian pos + D3PM type) score model.

Functions¶

`categorical_kl`(log_prob1, log_prob2)
`cosine_beta_schedule`(→ numpy.ndarray)
`extract`(coef, t, batch)
`get_beta_schedule`(→ numpy.ndarray)
`index_to_log_onehot`(→ torch.Tensor)
`log_1_min_a`(→ numpy.ndarray)
`log_add_exp`(→ torch.Tensor)
`log_categorical`(log_x_start, log_prob)
`log_normal`(values, means, log_scales)
`log_sample_categorical`(→ torch.Tensor)
`normal_kl`(mean1, logvar1, mean2, logvar2)
`threshold_cfg`(→ torch.Tensor)
`to_torch_const`(→ torch.nn.Parameter)

Module Contents¶

class MolecularDiffusion.modules.models.diffsmol.score_model.ScorePosNet3D(config: Dict[str, Any], atom_vocab: list[str])¶

Bases: torch.nn.Module

Shape-conditioned hybrid (Gaussian pos + D3PM type) score model.

compute_v_Lt(log_v_model_prob, log_v0, log_v_true_prob, t, batch)¶

forward(ligand_pos_perturbed: torch.Tensor, ligand_v_perturbed: torch.Tensor, batch_ligand: torch.Tensor, ligand_shape: torch.Tensor, time_step: torch.Tensor | None = None, return_all: bool = False) → Dict[str, torch.Tensor]¶

Predict (x0, v0) from the noisy sample at step t.

ligand_shape: (B, shape_dim, 3) equivariant latent, in the same frame as ligand_pos_perturbed.

get_diffusion_loss(ligand_pos: torch.Tensor, ligand_v: torch.Tensor, batch_ligand: torch.Tensor, ligand_shape: torch.Tensor, time_step: torch.Tensor | None = None, eval_mode: bool = False) → Dict[str, torch.Tensor]¶

Bond-free loss_pos + loss_v_weight * loss_v.

ligand_pos must already be surface-centered; no re-centering is applied (center_pos_mode: none upstream).

q_pos_posterior(x0, xt, t, batch)¶

q_v_posterior(log_v0, log_vt, t, batch)¶: q(v_{t-1} | v_t, v_0).

q_v_pred(log_v0, t, batch)¶: q(v_t | v_0).

q_v_pred_one_timestep(log_vt_1, t, batch)¶: q(v_t | v_{t-1}).

q_v_sample(log_v0, t, batch, num_classes)¶

sample_diffusion(init_ligand_pos: torch.Tensor, init_ligand_v: torch.Tensor, batch_ligand: torch.Tensor, ligand_shape: torch.Tensor, num_steps: int | None = None, guide_stren: float = 0.0, threshold_type: str | None = None, threshold_p: float | None = None, progress: bool = False) → Dict[str, torch.Tensor]¶

Reverse diffusion with classifier-free guidance.

guide_stren is the CFG weight w: the guided prediction is uncond + (1 + w) * (cond - uncond). 0 means conditional-only (no second pass), matching upstream’s shipped no-guidance setting.

sample_time(num_graphs: int, device)¶

alphas_cumprod¶

alphas_cumprod_prev¶

atom_vocab¶

betas¶

cond_mask_prob¶

config¶

ligand_atom_emb¶

log_alphas_cumprod_v¶

log_alphas_v¶

log_one_minus_alphas_cumprod_v¶

log_one_minus_alphas_v¶

loss_pos_type¶

loss_v_weight¶

loss_weight_type¶

model_mean_type¶

num_classes¶

num_timesteps¶

posterior_logvar¶

posterior_mean_c0_coef¶

posterior_mean_ct_coef¶

posterior_var¶

pred_bond_type = False¶

refine_net¶

sample_time_method¶

scalar_hidden_dim¶

shape_dim¶

sqrt_alphas_cumprod¶

sqrt_one_minus_alphas_cumprod¶

sqrt_recip_alphas_cumprod¶

sqrt_recipm1_alphas_cumprod¶

time_emb_dim¶

v_mode¶

v_net_type¶

vec_hidden_dim¶

MolecularDiffusion.modules.models.diffsmol.score_model.categorical_kl(log_prob1, log_prob2)¶

MolecularDiffusion.modules.models.diffsmol.score_model.cosine_beta_schedule(timesteps: int, s: float = 0.008) → numpy.ndarray¶

MolecularDiffusion.modules.models.diffsmol.score_model.extract(coef: torch.Tensor, t: torch.Tensor, batch: torch.Tensor)¶

MolecularDiffusion.modules.models.diffsmol.score_model.get_beta_schedule(beta_schedule: str, num_diffusion_timesteps: int, **kwargs: Any) → numpy.ndarray¶

MolecularDiffusion.modules.models.diffsmol.score_model.index_to_log_onehot(x: torch.Tensor, num_classes: int) → torch.Tensor¶

MolecularDiffusion.modules.models.diffsmol.score_model.log_1_min_a(a: numpy.ndarray) → numpy.ndarray¶

MolecularDiffusion.modules.models.diffsmol.score_model.log_add_exp(a: torch.Tensor, b: torch.Tensor) → torch.Tensor¶

MolecularDiffusion.modules.models.diffsmol.score_model.log_categorical(log_x_start, log_prob)¶

MolecularDiffusion.modules.models.diffsmol.score_model.log_normal(values, means, log_scales)¶

MolecularDiffusion.modules.models.diffsmol.score_model.log_sample_categorical(logits: torch.Tensor) → torch.Tensor¶

MolecularDiffusion.modules.models.diffsmol.score_model.normal_kl(mean1, logvar1, mean2, logvar2)¶

MolecularDiffusion.modules.models.diffsmol.score_model.threshold_cfg(x0: torch.Tensor, x0_cond: torch.Tensor, threshold_type: str | None, p: float | None = None) → torch.Tensor¶

MolecularDiffusion.modules.models.diffsmol.score_model.to_torch_const(x: numpy.ndarray) → torch.nn.Parameter¶

MolecularDiffusion.modules.models.diffsmol.score_model.DEFAULT_CONFIG: Dict[str, Any]¶