MolecularDiffusion.modules.models.syncogen.diffusion.sampling.integrators.euler

Euler integrator for flow matching.

Classes

EulerIntegrator

Simple Euler integrator for flow matching.

Module Contents

class MolecularDiffusion.modules.models.syncogen.diffusion.sampling.integrators.euler.EulerIntegrator(inference_annealing: bool = False, annealing_coef: float = 1.0)

Bases: MolecularDiffusion.modules.models.syncogen.diffusion.sampling.integrators.base.IntegratorBase

Simple Euler integrator for flow matching.

Updates coordinates using: C_{t-dt} = C_t + v_theta(C_t, t) * dt where v_theta is the velocity field pointing from C_t toward C_0.

Initialize Euler integrator.

Parameters:

  • inference_annealing – Whether to apply inference annealing

  • annealing_coef – Coefficient for inference annealing

compute_velocity(C: torch.Tensor, C0_pred: torch.Tensor, t: torch.Tensor) → torch.Tensor

Compute the velocity field v_theta at time t.

For flow matching: v_t = (C_0 - C_t) / t This points from C_t toward C_0.

Parameters:

  • C – Current noisy coordinates

  • C0_pred – Predicted clean coordinates

  • t – Current timestep

Returns:

Velocity tensor

step(coords: MolecularDiffusion.modules.models.syncogen.api.atomics.coordinates.Coordinates, C0_pred: torch.Tensor, t: torch.Tensor, dt: float) → torch.Tensor

Perform one Euler integration step.

Parameters:

  • coords – Coordinates object with current noisy state

  • C0_pred – Predicted clean coordinates

  • t – Current timestep

  • dt – Time step size

Returns:

Updated coordinates tensor

Return type:

C_next

annealing_coef = 1.0
inference_annealing = False