Update beta_binom_post_plot.py

scripts/beta_binom_post_plot.py

@@ -1,60 +1,79 @@
 import superimport
-
-import numpy as np
+import jax.numpy as jnp
 import matplotlib.pyplot as plt
 import pyprobml_utils as pml
+from jax.scipy.stats import beta, bernoulli
 
-from scipy.stats import dirichlet
-
-#Points where we evaluate the pdf
-x = np.linspace(0.001, .999, 100)
+# Points where we evaluate the pdf
+x = jnp.linspace(0.001, 0.999, 100)
 
-#Given an alpha parameter, this returns a pdf function
-def MakeBeta(alpha):
-    def Beta(y):
-        return dirichlet.pdf([y, 1 - y], alpha)
-    Beta = np.vectorize(Beta)
-    return Beta
 
-#Makes strings for the legend:
-def MakeLabel(Data,which):
-    alpha = Data[which]
-    lab = which + " Be(" + str(alpha[0]) + ", " + str(alpha[1]) + ")"
-    return lab
-
-#Forms graph give the parameters of the prior, likelihood and posterior:
-def MakeGraph(Data,SaveName):
-    prior = MakeBeta(Data['prior'])(x)
-    likelihood = MakeBeta(Data['lik'])(x)
-    posterior = MakeBeta(Data['post'])(x)
+# Forms graph given the parameters of the prior, likelihood and posterior:
+def make_graph(data, save_name):
+    prior = beta.pdf(x, a=data["prior"]["a"], b=data["prior"]["b"])
+    n_0 = data["likelihood"]["n_0"]
+    n_1 = data["likelihood"]["n_1"]
+    samples = jnp.concatenate([jnp.zeros(n_0), jnp.ones(n_1)])
+    likelihood_function = jnp.vectorize(
+        lambda p: jnp.exp(bernoulli.logpmf(samples, p).sum())
+    )
+    likelihood = likelihood_function(x)
+    posterior = beta.pdf(x, a=data["posterior"]["a"], b=data["posterior"]["b"])
 
     fig, ax = plt.subplots()
-    ax.plot(x, prior, 'r', label=MakeLabel(Data, "prior"), linewidth=2.0)
-    ax.plot(x, likelihood, 'k--', label=MakeLabel(Data, "lik"), linewidth=2.0)
-    ax.plot(x, posterior, 'b--', label=MakeLabel(Data, "post"), linewidth=2.0)
-    ax.legend(loc='upper left', shadow=True)
-    pml.savefig(SaveName)
+    axt = ax.twinx()
+    fig1 = ax.plot(
+        x,
+        prior,
+        "k",
+        label=f"prior Beta({data['prior']['a']}, {data['prior']['b']})",
+        linewidth=2.0,
+    )
+    fig2 = axt.plot(x, likelihood, "r:", label=f"likelihood Bernoulli", linewidth=2.0)
+    fig3 = ax.plot(
+        x,
+        posterior,
+        "b-.",
+        label=f"posterior Beta({data['posterior']['a']}, {data['posterior']['b']})",
+        linewidth=2.0,
+    )
+    fig_list = fig1 + fig2 + fig3
+    labels = [fig.get_label() for fig in fig_list]
+    ax.legend(fig_list, labels, loc="upper left", shadow=True)
+    axt.set_ylabel("Likelihood")
+    ax.set_ylabel("Prior/Posterior")
+    ax.set_title(f"$N_0$:{n_0}, $N_1$:{n_1}")
+    pml.savefig(save_name)
     plt.show()
 
-Data1 = {'prior': [1, 1],
-       'lik': [5, 2],
-       'post': [5, 2]}
 
-Data2 = {'prior': [1, 1],
-       'lik': [41, 11],
-       'post': [41, 11]}
+data1 = {
+    "prior": {"a": 1, "b": 1},
+    "likelihood": {"n_0": 1, "n_1": 4},
+    "posterior": {"a": 5, "b": 2},
+}
+
+data2 = {
+    "prior": {"a": 1, "b": 1},
+    "likelihood": {"n_0": 10, "n_1": 40},
+    "posterior": {"a": 41, "b": 11},
+}
 
-Data3 = {'prior': [2, 2],
-       'lik': [5, 2],
-       'post': [6, 3]}
+data3 = {
+    "prior": {"a": 2, "b": 2},
+    "likelihood": {"n_0": 1, "n_1": 4},
+    "posterior": {"a": 6, "b": 3},
+}
 
-Data4 = {'prior': [2, 2],
-       'lik': [41, 11],
-       'post': [42, 12]}
+data4 = {
+    "prior": {"a": 2, "b": 2},
+    "likelihood": {"n_0": 10, "n_1": 40},
+    "posterior": {"a": 42, "b": 12},
+}
 
-MakeGraph(Data1, "betaPost1")
-MakeGraph(Data2, "betaPost2")
-MakeGraph(Data3, "betaPost3")
-MakeGraph(Data4, "betaPost4")
+make_graph(data1, "betaPost1")
+make_graph(data2, "betaPost2")
+make_graph(data3, "betaPost3")
+make_graph(data4, "betaPost4")
 
 plt.show()