MolecularDiffusion.modules.models.loqi.interpolant¶

LoQI coordinate interpolants: VDM diffusion and continuous flow matching.

Ported from others/LoQI/src/megalodon/interpolant/ (NVIDIA, Apache-2.0), narrowed to the two paths both LoQI configs actually take:

ContinuousDiffusionInterpolant with diffusion_type: vdm – the discrete-time, cosine-adaptive VDM schedule of loqi.yaml. The ddpm branch is not ported (no LoQI config selects it).
ContinuousFlowMatchingInterpolant with a linear continuous-time schedule, prediction_type: velocity and optimal_transport: 'rigid' – loqi_flow.yaml. The vpe schedule and the permutation OT branch are ported because they cost three lines and share code paths, but no LoQI config reaches them.

The discrete interpolants are deliberately absent: LoQI marks h, edge_attr and charges as discrete_null, i.e. supplied un-noised as conditioning, so nothing here ever noises a categorical variable.

Buffer names are load-bearing – they are the interpolants.x.* keys of the released checkpoints.

Classes¶

`ContinuousDiffusionInterpolant`	VDM continuous Gaussian diffusion over coordinates, discrete time.
`ContinuousFlowMatchingInterpolant`	Linear continuous-time flow matching with velocity prediction + rigid OT.
`CosineSchedule`	MiDi-style adaptive cosine schedule (`nu` reshapes the time axis).
`Interpolant`
`LinearSchedule`

Functions¶

`align_prior`(→ torch.Tensor)
`build_interpolant`(interpolant_type, **kwargs)	The two coordinate interpolants LoQI selects, by config string.
`build_scheduler`(scheduler_type, num_diffusion_timesteps)
`rigid_alignment`(→ torch.Tensor)	Kabsch-align `x_0` onto `x_1`; reflection-safe (`det(R) = 1`).

Module Contents¶

class MolecularDiffusion.modules.models.loqi.interpolant.ContinuousDiffusionInterpolant(prior_type: str = 'gaussian', diffusion_type: str = 'vdm', solver_type: str = 'sde', timesteps: int = 500, time_type: str = 'discrete', num_classes: int = 3, scheduler_type: str = 'cosine_adaptive', s: float = 0.008, sqrt: bool = False, nu: float = 1.0, clip: bool = True, com_free: bool = True, cut: bool = False)¶

Bases: Interpolant

VDM continuous Gaussian diffusion over coordinates, discrete time.

forward_schedule(batch, time)¶

interpolate(batch, x1, time)¶

loss_weight_t(time)¶

prior(batch, shape, device, x1=None)¶

reverse_schedule(batch, time)¶

snr(time)¶

step(batch, xt, x_hat, x0=None, time=None, dt=None)¶

com_free = True¶

diffusion_type = 'vdm'¶

num_classes = 3¶

class MolecularDiffusion.modules.models.loqi.interpolant.ContinuousFlowMatchingInterpolant(prior_type: str = 'gaussian', vector_field_type: str = 'standard', solver_type: str = 'ode', timesteps: int = 500, min_t: float = 0.01, time_type: str = 'continuous', num_classes: int = 3, scheduler_type: str = 'linear', s: float = 0.008, sqrt: bool = False, nu: float = 1.0, clip: bool = True, com_free: bool = True, noise_sigma: float = 0.0, optimal_transport: str | None = None, clip_t: float = 0.9, loss_weight_type: str = 'uniform', loss_t_scale: float = 0.1, inference_noise_sigma: float | None = None, prediction_type: str = 'data')¶

Bases: Interpolant

Linear continuous-time flow matching with velocity prediction + rigid OT.

equivariant_ot_prior(batch, data_chunk, permutation: bool = True)¶

forward_schedule(batch, time)¶

interpolate(batch, x1, time)¶

loss_weight_t(time)¶: None means “no per-molecule weighting”; the loss function treats a None batch weight as 1 (upstream InterpolantLossFunction line 175). loss_weight_type defaults to 'standard' through upstream’s builder, which falls through every branch and returns None – reproduced here rather than silently changed.

prior(batch, shape, device, x1=None)¶

prior_func(batch, shape, device, x1=None)¶

step(batch, xt, x_hat, x0=None, time=None, dt=None)¶

update_weight(t)¶

vector_field(batch, x1, xt, time)¶

clip_t = 0.9¶

com_free = True¶

inference_noise_sigma¶

loss_t_scale = 0.1¶

loss_weight_type = 'uniform'¶

max_t = 0.99¶

min_t = 0.01¶

noise_sigma = 0.0¶

num_classes = 3¶

optimal_transport = None¶

prediction_type = 'data'¶

schedule_type = 'linear'¶

vector_field_type = 'standard'¶

class MolecularDiffusion.modules.models.loqi.interpolant.CosineSchedule(num_diffusion_timesteps: int, s: float = 0.008, sqrt: bool = False, nu: float = 1.0, clip: bool = True, cut: bool = False, **kwargs)¶

Bases: _Schedule

MiDi-style adaptive cosine schedule (nu reshapes the time axis).

compute_alphas() → torch.Tensor¶

clip = True¶

cut = False¶

nu = 1.0¶

num_diffusion_timesteps¶

s = 0.008¶

sqrt = False¶

class MolecularDiffusion.modules.models.loqi.interpolant.Interpolant(prior_type: str, solver_type: str = 'sde', timesteps: int = 500, time_type: str = 'discrete')¶

Bases: torch.nn.Module

sample_time(num_samples: int, method: str = 'uniform', device: str | torch.device = 'cpu', mean: float = 0.0, scale: float = 0.81, min_t: float = 0.0) → torch.Tensor¶: Only uniform is ported – both LoQI configs use it.

prior_type¶

solver_type = 'sde'¶

time_type = 'discrete'¶

timesteps = 500¶

class MolecularDiffusion.modules.models.loqi.interpolant.LinearSchedule(num_diffusion_timesteps: int, **kwargs)¶

Bases: _Schedule

compute_alphas() → torch.Tensor¶

num_diffusion_timesteps¶

MolecularDiffusion.modules.models.loqi.interpolant.align_prior(prior_feat: torch.Tensor, dst_feat: torch.Tensor, permutation: bool = False, rigid_body: bool = False, n_alignments: int = 1) → torch.Tensor¶

MolecularDiffusion.modules.models.loqi.interpolant.build_interpolant(interpolant_type: str, **kwargs)¶: The two coordinate interpolants LoQI selects, by config string.

MolecularDiffusion.modules.models.loqi.interpolant.build_scheduler(scheduler_type: str, num_diffusion_timesteps: int, s: float = 0.008, sqrt: bool = False, nu: float = 1.0, clip: bool = True, cut: bool = True)¶

MolecularDiffusion.modules.models.loqi.interpolant.rigid_alignment(x_0: torch.Tensor, x_1: torch.Tensor) → torch.Tensor¶: Kabsch-align x_0 onto x_1; reflection-safe (det(R) = 1).