MolecularDiffusion.modules.models.reactot.schedule¶

The Schrodinger-bridge schedule React-OT replaces DDPM’s noise with.

Vendored from reactot/diffusion/_schedule.py and reactot/diffusion/_utils.py of https://github.com/deepprinciple/react-ot at commit 6dfccd0.

SBSchedule is a plain Python class, not an nn.Module. It holds no buffers and nothing of it reaches a state dict – which is exactly why React-OT’s released checkpoint has 246 tensors where OA-ReactDiff’s has 248 (OA carries two gamma buffers of 5001 + 151 values, and 10,651,063 - 10,645,911 = 5,152).

A load-bearing degeneracy in the released settings¶

At timesteps=3000 / beta_max=0.3 / power=0.5 the schedule is constant:

linear_end = beta_max / timesteps = 1e-4 = make_beta_schedule's own
hardcoded linear_start, so the linspace is degenerate; ** 0.5 keeps it
constant; the mirror-concatenate keeps it constant; and the final
renormalisation lands every entry on 5e-5.

ode_sampling() asserts that constancy (en_sb.py:411 upstream). Changing timesteps or beta_max therefore breaks the ODE solver with an assertion, and changes the DDPM solver’s meaning silently. Both must stay at 3000 / 0.3; see the plan’s Hyperparameter Provenance table.

Classes¶

SBSchedule

The bridge's variance schedule.

Functions¶

`compute_gaussian_product_coef`(→ Tuple[Any, Any, Any])	Coefficients of the product of two Gaussians.
`make_beta_schedule`(→ numpy.ndarray)	Betas for the bridge, as upstream builds them.
`space_indices`(→ List[int])	`count` indices spread evenly over `range(num_steps)`.
`unsqueeze_xdim`(→ torch.Tensor)	Append `len(xdim)` trailing singleton axes to `z`.

Module Contents¶

class MolecularDiffusion.modules.models.reactot.schedule.SBSchedule(timesteps: int = 1000, beta_max: float = 0.3, power: float = 1.0, inv_power: float = 1.0)¶

The bridge’s variance schedule.

Not an nn.Module: see the module docstring. Everything is derived from betas, so a task rebuilds it from four numbers at construction and the checkpoint carries none of it.

timesteps¶: grid length, taken from the built array rather than the argument (the mirror-concatenate makes them equal for even timesteps, which 3000 is).

betas¶: (T,) per-step variance increments.

std_fwd¶: sqrt(cumsum(betas)).

std_bwd¶: sqrt(reverse cumsum(betas)).

std_sb¶: sqrt(var) of the bridge marginal.

mu_x0¶: weight on the endpoint x0 in q(x_t | x_0, x_1).

mu_x1¶: weight on the endpoint x1.

Build the schedule.

Parameters:

timesteps – grid length. 3000 for the released weights.
beta_max – total variance. 0.3 for the released weights.
power – exponent on the linspace.
inv_power – exponent on its endpoints.

get_std_fwd(step: torch.Tensor, xdim: Any = None) → torch.Tensor¶

std_fwd at step, optionally broadcast to xdim.

Parameters:

step – integer index tensor.
xdim – trailing shape to broadcast against; None => none.

Returns:

The looked-up standard deviations.

static inflate_batch_array(array: torch.Tensor, target: torch.Tensor) → torch.Tensor¶

Reshape a per-node vector to broadcast against target.

Parameters:

array – (n,) (or (n, 1, ..., 1)).
target – the tensor whose rank it must match.

Returns:

array viewed as (n, 1, ..., 1).

betas¶

mu_x0¶

mu_x1¶

std_bwd¶

std_fwd¶

std_sb¶

timesteps¶

MolecularDiffusion.modules.models.reactot.schedule.compute_gaussian_product_coef(sigma1: Any, sigma2: Any) → Tuple[Any, Any, Any]¶

Coefficients of the product of two Gaussians.

Given p1 = N(x_t | x_0, sigma1**2) and p2 = N(x_t | x_1, sigma2**2), return (coef1, coef2, var) such that p1 * p2 = N(x_t | coef1 * x0 + coef2 * x1, var).

Deliberately untyped in its operands: it is called both with numpy arrays (schedule construction) and with torch tensors (EnSB.p_posterior()), and does nothing either cannot do.

Parameters:

sigma1 – first standard deviation.
sigma2 – second standard deviation.

Returns:

(coef1, coef2, var).

MolecularDiffusion.modules.models.reactot.schedule.make_beta_schedule(n_timestep: int = 1000, linear_start: float = 0.0001, linear_end: float = 0.02, power: float = 1.0, inv_power: float = 1.0) → numpy.ndarray¶

Betas for the bridge, as upstream builds them.

Parameters:

n_timestep – length of the grid.
linear_start – first endpoint, before inv_power / power.
linear_end – last endpoint; the caller passes beta_max/timesteps.
power – exponent applied to the whole linspace.
inv_power – exponent applied to each endpoint first.

Returns:

(n_timestep,) float64 array.

MolecularDiffusion.modules.models.reactot.schedule.space_indices(num_steps: int, count: int) → List[int]¶

count indices spread evenly over range(num_steps).

This is what turns React-OT’s 3000-step grid into an nfe-step sampling schedule: space_indices(3000, nfe + 1) picks the nfe + 1 grid points the sampler walks backwards through, so the number of network evaluations is nfe.

Parameters:

num_steps – size of the grid to sample from.
count – how many indices to take.

Returns:

Ascending list of count indices, starting at 0.

MolecularDiffusion.modules.models.reactot.schedule.unsqueeze_xdim(z: torch.Tensor, xdim: Sequence[int]) → torch.Tensor¶

Append len(xdim) trailing singleton axes to z.

Parameters:

z – tensor to broadcast.
xdim – the trailing shape it must broadcast against.

Returns:

z viewed with the extra axes.