MolecularDiffusion.modules.models.chefnmr.diffusion¶

References: Karras et al. 2022 (EDM); Abramson et al. 2024 (AlphaFold3).

Ported unchanged except for two things that are packaging, not maths:

score_model_args is the denoiser’s kwargs dict directly, instead of upstream’s {model_name: str, <model_name>: {...}} indirection – there is exactly one score model and Hydra already picks the task.
einops.rearrange(sigma, "b -> b 1 1") is sigma[:, None, None]. einops is not installed here and this was its only site in this file.

sigma_data enters the preconditioning (c_in/c_sigma/d_sigma), so it is baked into the weights: a checkpoint must be sampled with the same value it was trained with, which is why it is a task-config key rather than something derived from whatever dataset happens to be attached.

Classes¶

AtomDiffusion

Forward noising + reverse (Heun/EDM) sampling over atom coordinates.

Module Contents¶

class MolecularDiffusion.modules.models.chefnmr.diffusion.AtomDiffusion(score_model_args: dict, train_sigma_distribution_type: str = 'af3', sample_sigma_schedule_type: str = 'edm', sample_gamma_schedule_type: str = 'edm', num_sampling_steps: int = 50, sigma_min: float = 0.0004, sigma_max: float = 80.0, gamma_min: float = 1.0, noise_scale: float = 1.0, step_scale: float = 1.0, guidance_scale: float = 0.0, synchronize_sigmas: bool = False, coordinate_transformation_when_training: str = 'centering_rotation_translation', edm_args: Dict | None = None, train_sigma_args: Dict | None = None, **kwargs)¶

Bases: torch.nn.Module

Forward noising + reverse (Heun/EDM) sampling over atom coordinates.

a_sigma(sigma: torch.Tensor) → torch.Tensor¶

b_sigma(sigma: torch.Tensor) → torch.Tensor¶

c_in(sigma: torch.Tensor) → torch.Tensor¶

c_sigma(sigma: torch.Tensor) → torch.Tensor¶

compute_loss(model_inputs: Dict[str, torch.Tensor], dict_out: Dict[str, torch.Tensor], multiplicity: int = 1, add_smooth_lddt_loss: bool = True, lddt_loss_threshold: list | None = None) → Dict[str, Any]¶

d_sigma(sigma: torch.Tensor) → torch.Tensor¶

float_to_tensor(value: float | torch.Tensor, batch_size: int, device) → torch.Tensor¶

forward(model_inputs: Dict[str, torch.Tensor], atom_coords: torch.Tensor, multiplicity: int = 1) → Dict[str, torch.Tensor]¶

interpolate(atom_coords, noise, sigma)¶

net_target(atom_coords, noise, sigma)¶

neural_network_forward(noisy_atom_coords: torch.Tensor, sigma: float | torch.Tensor, network_condition_kwargs: dict) → Dict[str, torch.Tensor]¶

noised_coords_in_network(atom_coords, sigma)¶

pad_sigma(sigma: torch.Tensor, batch_size: int, device) → torch.Tensor¶

predict_denoised_atom_coords(noisy_atom_coords, net_out, sigma)¶

predict_velocity(noisy_atom_coords, net_out, sigma, denoised_atom_coords=None)¶

sample(model_inputs: Dict[str, torch.Tensor], num_sampling_steps: int | None = None, multiplicity: int = 1, n_chain_frames: int = 1, guidance_scale: float | None = None) → Tuple[torch.Tensor, torch.Tensor]¶

Reverse diffusion.

guidance_scale=None uses self.guidance_scale (the model’s own default). It is an explicit argument rather than a mutated attribute so a generator can vary it per run without touching the module.

sample_gamma_schedule(sigmas: torch.Tensor) → torch.Tensor¶

sample_sigma_schedule(num_sampling_steps=None) → torch.Tensor¶

sigma_in_network(sigma: torch.Tensor) → torch.Tensor¶

train_sigma_distribution(batch_size: int) → torch.Tensor¶

coordinate_transformation_when_training = 'centering_rotation_translation'¶

property device: torch.device¶

edm_args¶

gamma_min = 1.0¶

guidance_scale = 0.0¶

noise_scale = 1.0¶

num_sampling_steps = 50¶

sample_gamma_schedule_type = 'edm'¶

sample_sigma_schedule_type = 'edm'¶

score_model¶

sigma_max = 80.0¶

sigma_min = 0.0004¶

step_scale = 1.0¶

synchronize_sigmas = False¶

train_sigma_args¶

train_sigma_distribution_type = 'af3'¶