MolecularDiffusion.modules.models.syncogen.core

Classes

Diffusion

Functions

get_backbone(backbone, self_conditioning, ...)

Module Contents

class MolecularDiffusion.modules.models.syncogen.core.Diffusion(data_manager: Any, losses: Sequence[MolecularDiffusion.modules.models.syncogen.diffusion.loss.LossBase] = (), device: Optional[torch.device] = torch.device('cuda' if torch.cuda.is_available() else 'cpu'), augmentations: Optional[list[str]] = ['center', 'normalize', 'random_rotate'], normalization_scale: Optional[float] = 1.0 / COORDS_STD, discrete_noise: Optional[MolecularDiffusion.modules.models.syncogen.diffusion.noise.NoiseBase] = LogLinearNoise(), interpolator: Optional[MolecularDiffusion.modules.models.syncogen.diffusion.interpolation.InterpolatorBase] = LinearInterpolator(), discrete_strategy: Optional[MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.DiscreteStrategyBase] = MDLM(), integrator: Optional[MolecularDiffusion.modules.models.syncogen.diffusion.sampling.integrators.IntegratorBase] = None, train_rot_align: bool = False, self_conditioning: bool = False, time_conditioning: bool = True, backbone: str = 'semla_pharm', use_compat: bool = True, sampling_eps: float = 0.001, importance_sampling: bool = False, antithetic_sampling: bool = True, sampling_noise_removal: bool = True, generate_eval_samples: bool = True, sample_every_n_epochs: int = 1, num_sample_steps: int = 100, ema_decay: float = 0.0, pharm_subset: int = 7, scale_noise: bool = False, scale_noise_factor: float = 0.2, num_fragments_probs: Optional[Dict[int, float]] = None, backbone_kwargs: Optional[Dict[str, Any]] = None)

Bases: torch.nn.Module

forward(Xt, Et, Ct, atom_mask, node_mask, cond=None, sampling=False)

Returns log score.

get_pharm_cond(batch)

Extract and augment pharmacophore conditioning from a batch.

restore_model_and_sample(num_steps: int, test_dataloader=None, cond: tuple | None = None)

Generate samples under EMA weights, restoring training weights afterwards.

Parameters:

  • num_steps – Number of diffusion steps

  • test_dataloader – Optional dataloader for extracting pharmacophores (training mode)

  • cond – Optional (types, pos, mask) tuple for direct pharmacophore conditioning (sampling mode)

sample(num_steps: int = 100, test_dataloader=None, cond: tuple | None = None)

Sample a batch of graphs from the model.

Parameters:

  • num_steps – Number of diffusion steps

  • test_dataloader – Optional dataloader for extracting pharmacophores (training mode)

  • cond – Optional (types, pos, mask) tuple for direct pharmacophore conditioning (sampling mode)

antithetic_sampling = True
augmentations = ['center', 'normalize', 'random_rotate']
backbone
property batch_size: int

Training batch size.

data_manager
property device: torch.device
discrete_noise
discrete_strategy
ema = None
property eval_batch_size: int

Evaluation/validation batch size.

fast_forward_batches = None
fast_forward_epochs = None
generate_eval_samples = True
importance_sampling = False
integrator
interpolator
losses
n_edge_features
neg_infinity = -25000.0
normalization_scale = 0.3576209484116885
num_sample_steps = 100
pharm_subset = 7
sample_every_n_epochs = 1
sampling_eps = 0.001
sampling_noise_removal = True
scale_noise = False
scale_noise_factor = 0.2
self_conditioning = False
time_conditioning = True
train_rot_align = False
use_compat = True
MolecularDiffusion.modules.models.syncogen.core.get_backbone(backbone: str, self_conditioning: bool, pharm_subset: int, **backbone_kwargs)