MolecularDiffusion.modules.models.flowmol.interpolant_scheduler¶

Interpolant schedule for FlowMol flow-matching.

Ported verbatim from FlowMol (flowmol/models/interpolant_scheduler.py). Governs the per-feature interpolation weights between prior (t=0) and data (t=1) and the time-dependent loss weights.

Classes¶

InterpolantScheduler

Module Contents¶

class MolecularDiffusion.modules.models.flowmol.interpolant_scheduler.InterpolantScheduler(canonical_feat_order: list, schedule_type: str | Dict[str, str] = 'cosine', cosine_params: dict = {})¶

Bases: torch.nn.Module

alpha_t(t: torch.Tensor) → torch.Tensor¶

alpha_t_prime(t: torch.Tensor) → torch.Tensor¶

cosine_alpha_t(t: torch.Tensor, nu: torch.Tensor) → torch.Tensor¶

cosine_alpha_t_prime(t: torch.Tensor, nu: torch.Tensor) → torch.Tensor¶

interpolant_weights(t: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]¶: Weights for x_0 and x_1 in the interpolation between them.

linear_alpha_t(t: torch.Tensor) → torch.Tensor¶

linear_alpha_t_prime(t: torch.Tensor) → torch.Tensor¶

loss_weights(t: torch.Tensor)¶

update_device(t)¶

clamp_t = True¶

cosine_params¶

feats¶

n_feats¶

schedule_types¶

supported_schedule_types = ['cosine', 'linear']¶