From 4e45e80051db09b4d84620249eb80ca40831041e Mon Sep 17 00:00:00 2001
From: Kevin Wu <t-kevinwu@microsoft.com>
Date: Thu, 11 Aug 2022 22:03:06 +0000
Subject: [PATCH] More comprehensive variance plots

---
 .../cosine_var_schedule.pdf                   | Bin 6059 -> 6059 bytes
 .../linear_var_schedule.pdf                   | Bin 6297 -> 6297 bytes
 .../quadratic_var_schedule.pdf                | Bin 7041 -> 7041 bytes
 protdiff/beta_schedules.py                    |   7 ++++---
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/plots/variance_schedules/cosine_var_schedule.pdf b/plots/variance_schedules/cosine_var_schedule.pdf
index a3d9f22a9f84d763e09c034fa6165d734413b176..3133894316c3491a601a8788cdceed6118fefdef 100644
GIT binary patch
delta 22
ecmZ3jzgmC8XAw>dLqj7&BO_D8&3{DJGXVfxr3bJ8

delta 22
ecmZ3jzgmC8XAw?wLt`T&V@ng0&3{DJGXVfy1qZwU

diff --git a/plots/variance_schedules/linear_var_schedule.pdf b/plots/variance_schedules/linear_var_schedule.pdf
index 2daa588ca4fa75cc09ae115df1742065e25ae289..70562cbfc1485cc41d0cf85eddd7211d5d2746d3 100644
GIT binary patch
delta 22
ecmbPfIMZ;$PEk$^Lqj7&BO_D8&4)y<GXVfqq6bC*

delta 22
ecmbPfIMZ;$PEk&CLt`T&V@ng0&4)y<GXVfr0tZq6

diff --git a/plots/variance_schedules/quadratic_var_schedule.pdf b/plots/variance_schedules/quadratic_var_schedule.pdf
index 8469237bb74439f63868b9840f9f01f128f95bff..b148fdfa63c24d17551a10f3b78ff1e57003f35d 100644
GIT binary patch
delta 22
dcmZoPZ#3VqP=eFK(9p=x$jH=i^D2q!OaM^h2KWE~

delta 22
dcmZoPZ#3VqP=eFk(Adbx*wVyg^D2q!OaM_s2Lu2B

diff --git a/protdiff/beta_schedules.py b/protdiff/beta_schedules.py
index ac3a9e5..0918e30 100644
--- 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__":