docstrings plot code

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -452,13 +452,14 @@ class ShiftedBetaGeometric(Discrete):
         import scipy.stats as st
         from scipy.special import beta
         import arviz as az
+
         plt.style.use('arviz-darkgrid')
         t = np.arange(1, 11)
-        alpha_vals = [1., 1., 2., 2.]
-        beta_vals = [.1, .25, .5, 1.]
+        alpha_vals = [.1, .1, 1, 1]
+        beta_vals = [.5, 1, 1, 4]
         for alpha, _beta in zip(alpha_vals, beta_vals):
-            pmf = beta(alpha + 1, beta + t - 1) / beta(alpha, beta)
-            plt.plot(t, pmf, '-o', label=r'$\alpha$ = {}, $beta$ = {}'.format(alpha, beta))
+            pmf = beta(alpha + 1, _beta + t - 1) / beta(alpha, _beta)
+            plt.plot(t, pmf, '-o', label=r'$\alpha$ = {}, $beta$ = {}'.format(alpha, _beta))
         plt.xlabel('t', fontsize=12)
         plt.ylabel('p(t)', fontsize=12)
         plt.legend(loc=1)