More comprehensive variance plots

2026-06-04 21:34:32 +08:00 · 2022-08-11 22:03:06 +00:00
parent fa67d995f0
commit 4e45e80051
4 changed files with 4 additions and 3 deletions
--- a/plots/variance_schedules/cosine_var_schedule.pdf
+++ b/plots/variance_schedules/cosine_var_schedule.pdf
--- a/plots/variance_schedules/linear_var_schedule.pdf
+++ b/plots/variance_schedules/linear_var_schedule.pdf
--- a/plots/variance_schedules/quadratic_var_schedule.pdf
+++ b/plots/variance_schedules/quadratic_var_schedule.pdf
--- a/protdiff/beta_schedules.py
+++ b/protdiff/beta_schedules.py
@@ -86,10 +86,11 @@ def plot_variance_schedule(
    logging.info(
        f"Plotting {keyword} variance schedule with {timesteps} timesteps, ranging from {np.min(variance_vals)}-{np.max(variance_vals)}"
    )
-
+    alpha_beta_vals = compute_alphas(variance_vals)
    fig, ax = plt.subplots(dpi=300)
-    ax.plot(np.arange(timesteps), variance_vals)
-    fig.savefig(fname)
+    for k, v in alpha_beta_vals.items():
+        ax.plot(np.arange(timesteps), v.numpy(), label=k, alpha=0.7)
+    fig.savefig(fname + "_alpha_beta.pdf")


 if __name__ == "__main__":