Diffusion from noise to data

import numpy as np

# The demo's schedule: beta(t) rises from 0.1 to 20 as t goes from 0 to 1.
def beta(t):
    return 0.1 + 19.9 * t

def alpha_bar(t):
    return np.exp(-(0.1 * t + 0.5 * 19.9 * t**2))

# Forward process: keep sqrt(abar) of a clean point, add sqrt(1 - abar) of noise.
x0, eps = np.array([1.5, 0.0]), np.array([0.3, -1.2])
for t in (0.1, 0.3, 1.0):
    ab = alpha_bar(t)
    xt = np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps
    print(f"t = {t}: signal kept {np.sqrt(ab):.3f}, x_t = ({xt[0]:.3f}, {xt[1]:.3f})")

# Data: two small Gaussian blobs. Noising them gives two wider, shrunken
# blobs, so the score (the gradient of log density) has an exact formula.
means, std = np.array([[-1.0, 0.0], [1.0, 0.0]]), 0.1

def score(x, t):
    ab = alpha_bar(t)
    m, v = np.sqrt(ab) * means, ab * std**2 + 1 - ab
    logits = -np.sum((x - m) ** 2, axis=1) / (2 * v)
    r = np.exp(logits - logits.max())
    r /= r.sum()  # probability that x came from each blob
    return r @ (m - x) / v

def drift(x, t):  # probability flow ODE: dx/dt = -beta(t)/2 (x + score)
    return -0.5 * beta(t) * (x + score(x, t))

# Run time backward from 1 to 0.001 in 140 Heun steps, as the demo does.
# There is no randomness: the same start always ends in the same place.
def sample(x):
    ts = 0.001 + 0.999 * np.linspace(1, 0, 141) ** 2
    for t0, t1 in zip(ts[:-1], ts[1:]):
        k1 = drift(x, t0)
        k2 = drift(x + (t1 - t0) * k1, t1)
        x = x + (t1 - t0) * (k1 + k2) / 2
    return x

for start in ([0.5, 0.3], [-0.2, 1.0], [0.5, 0.3]):
    end = sample(np.array(start))
    print(f"noise {start} -> sample ({end[0]:.3f}, {end[1]:.3f})")

t = 0.1: signal kept 0.947, x_t = (1.517, -0.386)
t = 0.3: signal kept 0.630, x_t = (1.177, -0.932)
t = 1.0: signal kept 0.007, x_t = (0.310, -1.200)
noise [0.5, 0.3] -> sample (0.970, 0.030)
noise [-0.2, 1.0] -> sample (-0.899, 0.101)
noise [0.5, 0.3] -> sample (0.970, 0.030)