MolecularDiffusion.modules.models.shepherd_arch.inference.noise

This module contains the noise functions for the inference sampler.

Functions

forward_jump(→ tuple[torch.Tensor, int])

Perform a forward jump for a given timestep.

forward_jump_parameters(t_start_idx, jump, sigma_ts)

Compute the forward jump parameters for a given timestep.

forward_trajectory(→ dict[int, torch.Tensor])

Simulate a forward noising trajectory.

subsample_noise_schedule(num_steps, noise_schedule)

Subsample the noise schedule to a given number of steps.

Module Contents

MolecularDiffusion.modules.models.shepherd_arch.inference.noise.forward_jump(x: torch.Tensor, t_start, jump, sigma_ts, remove_COM_from_noise=False, batch=None, mask=None) → tuple[torch.Tensor, int]

Perform a forward jump for a given timestep.

Parameters:

  • x (The input tensor.)

  • t_start (The starting timestep.)

  • jump (The number of timesteps to jump.)

MolecularDiffusion.modules.models.shepherd_arch.inference.noise.forward_jump_parameters(t_start_idx, jump, sigma_ts)

Compute the forward jump parameters for a given timestep.

Parameters:

  • t_start_idx (The starting index of the timestep.)

  • jump (The number of timesteps to jump.)

  • sigma_ts (The standard deviation schedule.)

MolecularDiffusion.modules.models.shepherd_arch.inference.noise.forward_trajectory(x, ts, alpha_ts, sigma_ts, remove_COM_from_noise=False, mask=None, deterministic=False, batch=None, batch_size: int | None = None) → dict[int, torch.Tensor]

Simulate a forward noising trajectory.

Parameters:

  • x (The input tensor.)

  • ts (The timesteps.)

  • alpha_ts (The alpha schedule.)

  • sigma_ts (The sigma schedule.)

  • remove_COM_from_noise (Whether to remove the center of mass from the noise.)

  • mask (The mask tensor.)

  • deterministic (Whether to use a deterministic trajectory)

  • (i.e.

  • trajectory). (clean structure for entire)

  • batch (The batch tensor.)

  • batch_size (If provided and > 1, the input x (and corresponding mask/batch if given)) – will be repeated batch_size times to create batch_size independent trajectories.

MolecularDiffusion.modules.models.shepherd_arch.inference.noise.subsample_noise_schedule(num_steps: int, noise_schedule: dict)

Subsample the noise schedule to a given number of steps. Keep the last step of the original schedule.

Parameters:

  • num_steps (The number of steps to subsample to.)

  • noise_schedule (The noise schedule to subsample.)

Returns:

noise_schedule

Return type:

The subsampled noise schedule.