import numpy as np
import matplotlib.pyplot as plt

class Scheduler:
    def __init__(self, beta_start=1e-4, beta_end=2e-2, num_train_timesteps=1000):
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.num_train_timesteps = num_train_timesteps

        self.betas = np.linspace(self.beta_start, self.beta_end, self.num_train_timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = np.cumprod(self.alphas, axis=0)

sch = Scheduler(beta_start=1e-4, beta_end=2e-2, num_train_timesteps=1000)

plt.plot([sch.alphas[t] for t in range(sch.num_train_timesteps)], label=r"$\alpha_t = 1 - \beta_t$")
plt.plot([sch.alphas_cumprod[t] for t in range(sch.num_train_timesteps)], label=r"$\bar\alpha_t$")
plt.plot([np.sqrt(1 - sch.alphas_cumprod[t]) for t in range(sch.num_train_timesteps)], label=r"$\sigma_t = \sqrt[[:Template:1 - \bar\alpha t]]$")
plt.plot([sch.betas[t] for t in range(sch.num_train_timesteps)], label=r"$\beta_t$")
plt.legend()
plt.yscale('log')
plt.xlabel(f'$t