forward.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import beta, norm

# Base distribution: random Beta(k1, k2)
k1, k2 = np.random.uniform(0.5, 5), np.random.uniform(0.5, 5)
x0 = beta.rvs(k1, k2, size=10000)

# Diffusion parameters
beta_t = 0.05  # noise scale
timesteps = np.linspace(0, 100, 25)  # 6 steps including t=0

# Plot
plt.figure(figsize=(10, 5))
for t in timesteps:
    alpha_bar = (1 - beta_t) ** t
    eps = norm.rvs(size=x0.shape)
    xt = np.sqrt(alpha_bar) * x0 + np.sqrt(1 - alpha_bar) * eps
    plt.hist(xt, bins=80, density=True, alpha=0.5, label=f"t={t:.2f}")

plt.title("Diffusion Forward Process on Random Beta Distribution")
plt.xlabel("x_t")
plt.ylabel("Density")
plt.legend()
plt.grid(True)
plt.show()