<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Dice, Polls & Dirichlet Multinomials
As part of a longer term project to learn Bayesian Statistics, I'm currently reading Bayesian Data Analysis, 3rd Edition by Andrew Gelman, John Carlin, Hal Stern, David Dunson, Aki Vehtari, and Donald Rubin, commonly known as BDA3.
Although I've been using Bayesian statistics and probabilistic programming languages, like PyMC3, in projects for the last year or so, this book forces me to go beyond a pure practioner's approach to modeling, while still delivering very practical value.
Below are a few take aways from the earlier chapters in the book I found interesting. They are meant to hopefully inspire others to learn about Bayesian statistics, without trying to be overly formal about the math. If something doesn't look 100% to the trained mathematicians in the room, please let me know, or just squint a little harder. ;)
We'll cover
Step1:  Just looking at a simple bar plot, we suspect that we might not be dealing with a fair die! 
However, students of Bayesian statistics that we are, we'd like to go further and quantify our uncertainty in the fairness of the die and calculate the probability that someone slipped us loaded dice.
Step2:  Let's set up a simple model in PyMC3 that not only calculates the posterior probability for $theta$ (i.e. the probability for each side of the die), but also estimates the bias for throwing a $6$. 
We will use Deterministic variable, in addition to our unobserved (theta) and observed (results) variables.
For the prior on $theta$, we'll use a non-informative uniform distribution, by initializing the $Dirichlet$ prior with a series of 1s for the parameter a, one for each of the k possible outcomes. This is similar to initializing a $Beta$ distribution as $Beta(1, 1)$, which corresponds to the Uniform distribution.
Step3:  Starting with version 3.5, PyMC3 includes a handy function to plot models in plate notation
Step4:  Let's draw 1,000 samples from the joint posterior using the default NUTS sampler
Step5:  From the traceplot, we can already see that one of the $theta$ posteriors isn't in line with the rest
Step6:  We'll plot the posterior distributions for each $theta$ and compare it our reference value $p$ to see if the 95% HPD (Highest Posterior Density) interval includes $p = 1/6$.
Step7:  We can clearly see that the HPD for the posterior probability for rolling a $6$ barely includes what we'd expect from a fair die.
To be more precise, let's plot the probability of our die being biased on $6$, by comparing $theta[Six]$ to $p$
Step8:  Lastly, we can calculate the probability that the die is biased on $6$ by calculating the density to the right of our reference line at $0$
Step9:  Better get some new dice...!
Polling #1
Let's turn our review of the Dirichlet-Multinomial distribution to another example, concerning polling data. 
In section 3.4 of BDA3 on multivariate models and, specifically the section on Multinomial Models for Categorical Data, the authors include a, little dated, example of polling data in the 1988 Presidential race between George H.W. Bush and Michael Dukakis.
Here's the setup
Step10:  We, again, set up a simple Dirichlet-Multinomial model and include a Deterministic variable that calculates the metric of interest - the difference in probability of respondents for Bush vs. Dukakis.
Step11:  Looking at the % difference between respondents for Bush vs Dukakis, we can see that most of the density is greater than 0%, signifying a strong advantage for Bush in this poll.
We've also fit a $Beta$ distribution to this data via scipy.stats, and we can see that posterior of the difference of the 2 $theta$ values is a pretty good match.
Step12:  Percentage of samples with bush_dukakis_diff &gt; 0
Step13:  Polling #2
As an extension to the previous model, the authors of BDA include an exercise in chapter 3.10 (Exercise 2) that presents us with polling data from the 1988 Presidential race, taking before and after the one of the debates.
Comparison of two multinomial observations
Step14:  Convert to 2x3 array
Step15:  Number of respondents in each survey
Step16:  Number of respondents for the 2 major candidates in each survey
Step17:  For this model, we'll need to set up the priors slightly differently. Instead of 1 set of thetas, we need 2, one for each survey (pre/post debate).
To do that without creating specific pre/post versions of each variable, we'll take advantage of PyMC3's shape parameter, available for most (all?) distributions.
In this case, we'll need a 2-dimensional shape parameter, representing the number of debates n_debates and the number of choices in candidates n_candidates
Step18:  Thus, we need to initialize a Dirichlet distribution prior with shape (2,3) and then refer to the relevant parameters by index where needed.
Step19:  For models with multi-dimensional shapes, it's always good to check the shapes of the various parameters before sampling
Step20:  The plate notation visual can also help with that
Step21:  Let's sample with a slightly higher number of draws and tuning steps
Step22:  We'll take a look at the means of the posteriors for theta, indicating the % of support for each candidate pre & post debate
Step23:  Just from the means, we can see that the number of Bush supporters has likely decreased post debate from 48.8% to 46.3% (as a % of supporters of the 2 major candidates)
Step24:  Let's compare the results visually, by plotting the posterior distributions of the pre/post debate values for % responses for Bush and the posterior for pre/post difference in Bush supporters
Step25:  From the second plot, we can already see that a large portion of the posterior density is below 0, but let's be precise and actually calculate the probability that support shifted towards Bush after the debate
<ASSISTANT_TASK:>
Python Code:
y = np.asarray([20,  21, 17, 19, 17, 28])
k = len(y)
p = 1/k
n = y.sum()
n, p
Explanation: Dice, Polls & Dirichlet Multinomials
As part of a longer term project to learn Bayesian Statistics, I'm currently reading Bayesian Data Analysis, 3rd Edition by Andrew Gelman, John Carlin, Hal Stern, David Dunson, Aki Vehtari, and Donald Rubin, commonly known as BDA3.
Although I've been using Bayesian statistics and probabilistic programming languages, like PyMC3, in projects for the last year or so, this book forces me to go beyond a pure practioner's approach to modeling, while still delivering very practical value.
Below are a few take aways from the earlier chapters in the book I found interesting. They are meant to hopefully inspire others to learn about Bayesian statistics, without trying to be overly formal about the math. If something doesn't look 100% to the trained mathematicians in the room, please let me know, or just squint a little harder. ;)
We'll cover:
- Some common conjugate distributions
- An example of the Dirichlet-Multinomial distribution using dice rolls
- Two examples involing polling data from BDA3
Conjugate Distributions
In Chapter 2 of the book, the authors introduce several choices for prior probability distributions, along with the concept of conjugate distributions in section 2.4. 
From Wikipedia
In Bayesian probability theory, if the posterior distributions p(θ | x) are in the same probability distribution family as the prior probability distribution p(θ), the prior and posterior are then called conjugate distributions, and the prior is called a conjugate prior for the likelihood function
John Cook has this helpful diagram on his website that shows some common families of conjugate distributions:
<img src=https://www.johndcook.com/conjugate_prior_diagram.png width="300">
Conjugate distributions are a very important concept in probability theory, owing to a large degree to some nice mathematical properties that make computing the posteriors more tractable. Even with increasingly better computational tools, such as MCMC, models based on conjugate distributions are advantageous.
Beta-Binomial
One of the better known examples of conjugate distributions is the Beta-Binomial distribution, which is often used to model series of coin flips (the ever present topic in posts about probability). While the $Binomial$ distribution represents the probability of success in a series of Bernoulli trials, the Beta distribution here represents the prior probability distribtution of the probability of success for each trial. 
Thus, the probability $p$ of a coin landing on head is modeled to be $Beta$ distributed (with parameters $\alpha$ and $\beta$), while the likelihood of heads and tails is assumed to follow a $Binomial$ distribution with parameters $n$ (representing the number of flips) and the $Beta$ distributed $p$, thus creating the link.
$$p \sim Beta(\alpha, \beta)$$
$$y \sim Binomial(n, p)$$
Gamma-Poisson
Another often-used conjugate distribution is the Gamma-Poisson distribution, so named because the rate parameter $\lambda$ that parameterizes the Poisson distributed is modeled as a Gamma distribution:
$$\lambda \sim Gamma(k, \theta)$$
$$y \sim Poisson(\lambda)$$
While the discrete $Poisson$ distributed is often used in applications of count data, such as store customers, eCommerce orders, website visits, the $Gamma$ distribution serves as a useful distribution to model the rate at which these events occur ($\lambda$), since the $Gamma$ distribution models positive continuous values only but is otherwise quite flexible:
<img src=https://upload.wikimedia.org/wikipedia/commons/e/e6/Gamma_distribution_pdf.svg width="500">
Dirichlet-Multinomial
A perhaps more interesting and seemingly less talked-about example of conjugate distributions is the Dirichlet-Multinomial distribution, introduced in chapter 3 of BDA3. 
One way of think about the $Dirichlet-Multinomial$ distribution is that while the $Multinomial$ (-> multiple choices) distribution is a generalization of the $Binomial$ distribution (-> binary choice), the $Dirichlet$ distribution is a generalization of the $Beta$ distribution. That is, while the $Beta$ distribution models the probability of a single probability $p$, the $Dirichlet$ models the probabilities of multiple, mutually exclusive choices, parameterized by $a$ which is referred to as the concentration parameter and represents the weights for each choice (we'll see more on that later). 
In other words, think coins for $Beta-Binomial$ and dice for $Dirichlet-Multinomial$.
$$\theta \sim Dirichlet(a)$$
$$y \sim Multinomial(n, \theta)$$
In the wild, we might encounter the Dirichlet distribution these days mostly in the context of topic modeling in natural language processing, where it's commonly used as part of a Latent Dirichlet Allocation (or LDA) model, which is fancy way of saying we're trying to figure out the probability of an article belonging to a certain topic given its text.
However, for our purposes, let's look at the Dirichlet-Multinomial in the context of multiple choices, and let's start by throwing dice as a motivating example:
Throwing Dice
Let's first create some data representing 122 rolls of six-sided die, where $p$ represents the expected probability for each side, $1/6$
End of explanation
sns.barplot(x=np.arange(1, k+1), y=y);
Explanation: Just looking at a simple bar plot, we suspect that we might not be dealing with a fair die! 
However, students of Bayesian statistics that we are, we'd like to go further and quantify our uncertainty in the fairness of the die and calculate the probability that someone slipped us loaded dice.
End of explanation
n, y
with pm.Model() as dice_model:
    
    # initializes the Dirichlet distribution with a uniform prior:
    a = np.ones(k) 
    
    theta = pm.Dirichlet("theta", a=a)
    
    # Since theta[5] will hold the posterior probability of rolling a 6
    # we'll compare this to the reference value p = 1/6
#     six_bias = pm.Deterministic("six_bias", theta[k-1] - p)
    
    results = pm.Multinomial("results", n=n, p=theta, observed=y)
dice_model
Explanation: Let's set up a simple model in PyMC3 that not only calculates the posterior probability for $theta$ (i.e. the probability for each side of the die), but also estimates the bias for throwing a $6$. 
We will use Deterministic variable, in addition to our unobserved (theta) and observed (results) variables.
For the prior on $theta$, we'll use a non-informative uniform distribution, by initializing the $Dirichlet$ prior with a series of 1s for the parameter a, one for each of the k possible outcomes. This is similar to initializing a $Beta$ distribution as $Beta(1, 1)$, which corresponds to the Uniform distribution.
End of explanation
pm.model_to_graphviz(dice_model)
dice_model.check_test_point()
Explanation: Starting with version 3.5, PyMC3 includes a handy function to plot models in plate notation:
End of explanation
with dice_model:
    dice_trace = pm.sample(draws=1000)
Explanation: Let's draw 1,000 samples from the joint posterior using the default NUTS sampler:
End of explanation
with dice_model:
    pm.traceplot(dice_trace, combined=True, lines={"theta": p})
Explanation: From the traceplot, we can already see that one of the $theta$ posteriors isn't in line with the rest:
End of explanation
axes = pm.plot_posterior(dice_trace, varnames=["theta"], ref_val=np.round(p, 3))
for i, ax in enumerate(axes):
    ax.set_title(f"{i+1}")
Explanation: We'll plot the posterior distributions for each $theta$ and compare it our reference value $p$ to see if the 95% HPD (Highest Posterior Density) interval includes $p = 1/6$.
End of explanation
ax = pm.plot_posterior(dice_trace, varnames=["six_bias"], ref_val=[0])
ax.set_title(f"P(Theta[Six] - {p:.2%})");
Explanation: We can clearly see that the HPD for the posterior probability for rolling a $6$ barely includes what we'd expect from a fair die.
To be more precise, let's plot the probability of our die being biased on $6$, by comparing $theta[Six]$ to $p$
End of explanation
six_bias_perc = len(dice_trace["six_bias"][dice_trace["six_bias"]>0])/len(dice_trace["six_bias"])
      
print(f'P(Six is biased) = {six_bias_perc:.2%}')
Explanation: Lastly, we can calculate the probability that the die is biased on $6$ by calculating the density to the right of our reference line at $0$:
End of explanation
y = np.asarray([727, 583, 137])
n = y.sum()
k = len(y)
n, k
Explanation: Better get some new dice...!
Polling #1
Let's turn our review of the Dirichlet-Multinomial distribution to another example, concerning polling data. 
In section 3.4 of BDA3 on multivariate models and, specifically the section on Multinomial Models for Categorical Data, the authors include a, little dated, example of polling data in the 1988 Presidential race between George H.W. Bush and Michael Dukakis.
Here's the setup:
1,447 likely voters were surveyed about their preferences in the upcoming presidential election
Their responses were:
Bush: 727
Dukakis: 583
Other: 137
What is the probability that more people will vote for Bush over Dukakis?
i.e. what is the difference in support for the two major candidates?
We set up the data, where $k$ represents the number of choices the respondents had:
End of explanation
with pm.Model() as polling_model:
    
    # initializes the Dirichlet distribution with a uniform prior:
    a = np.ones(k) 
    
    theta = pm.Dirichlet("theta", a=a)
    
    bush_dukakis_diff = pm.Deterministic("bush_dukakis_diff", theta[0] - theta[1])
    
    likelihood = pm.Multinomial("likelihood", n=n, p=theta, observed=y)
pm.model_to_graphviz(polling_model)
with polling_model:
    polling_trace = pm.sample(draws=1000)
with polling_model:
    pm.traceplot(polling_trace, combined=True)
Explanation: We, again, set up a simple Dirichlet-Multinomial model and include a Deterministic variable that calculates the metric of interest - the difference in probability of respondents for Bush vs. Dukakis.
End of explanation
_, ax = plt.subplots(1,1, figsize=(10, 6))
sns.distplot(polling_trace["bush_dukakis_diff"], bins=20, ax=ax, kde=False, fit=stats.beta)
ax.axvline(0, c='g', linestyle='dotted')
ax.set_title("% Difference Bush vs Dukakis")
ax.set_xlabel("% Difference");
Explanation: Looking at the % difference between respondents for Bush vs Dukakis, we can see that most of the density is greater than 0%, signifying a strong advantage for Bush in this poll.
We've also fit a $Beta$ distribution to this data via scipy.stats, and we can see that posterior of the difference of the 2 $theta$ values is a pretty good match.
End of explanation
bush_dukakis_diff_perc = len(polling_trace["bush_dukakis_diff"][polling_trace["bush_dukakis_diff"]>0])/len(polling_trace["bush_dukakis_diff"])
      
print(f'P(More Responses for Bush) = {bush_dukakis_diff_perc:.0%}')
Explanation: Percentage of samples with bush_dukakis_diff &gt; 0:
End of explanation
data = pd.DataFrame([
        {"candidate": "bush", "pre": 294, "post": 288},
        {"candidate": "dukakis", "pre": 307, "post": 332},
        {"candidate": "other", "pre": 38, "post": 10}
       ], columns=["candidate", "pre", "post"])
data
Explanation: Polling #2
As an extension to the previous model, the authors of BDA include an exercise in chapter 3.10 (Exercise 2) that presents us with polling data from the 1988 Presidential race, taking before and after the one of the debates.
Comparison of two multinomial observations: on September 25, 1988, the evening of a
presidential campaign debate, ABC News conducted a survey of registered voters in the
United States; 639 persons were polled before the debate, and 639 different persons were
polled after. The results are displayed in Table 3.2. Assume the surveys are independent
simple random samples from the population of registered voters. Model the data with
two different multinomial distributions. For $j = 1, 2$, let $\alpha_j$ be the proportion of voters
who preferred Bush, out of those who had a preference for either Bush or Dukakis at
the time of survey $j$. Plot a histogram of the posterior density for $\alpha_2 − \alpha_1$. What is the
posterior probability that there was a shift toward Bush?
Let's copy the data from the exercise and model the problem as a probabilistic model, again using PyMC3:
End of explanation
y = data[["pre", "post"]].T.values
y
Explanation: Convert to 2x3 array
End of explanation
n = y.sum(axis=1) 
n
Explanation: Number of respondents in each survey
End of explanation
m = y[:, :2].sum(axis=1) 
m
Explanation: Number of respondents for the 2 major candidates in each survey
End of explanation
n_debates, n_candidates = y.shape
n_debates, n_candidates
Explanation: For this model, we'll need to set up the priors slightly differently. Instead of 1 set of thetas, we need 2, one for each survey (pre/post debate).
To do that without creating specific pre/post versions of each variable, we'll take advantage of PyMC3's shape parameter, available for most (all?) distributions.
In this case, we'll need a 2-dimensional shape parameter, representing the number of debates n_debates and the number of choices in candidates n_candidates
End of explanation
with pm.Model() as polling_model_debates:
    
    # initializes the Dirichlet distribution with a uniform prior:
    shape = (n_debates, n_candidates)
    a = np.ones(shape)
    
    # This creates a separate Dirichlet distribution for each debate
    # where sum of probabilities across candidates = 100% for each debate
    theta = pm.Dirichlet("theta", a=a, shape=shape)
    
    # get the "Bush" theta for each debate, at index=0
    bush_pref = pm.Deterministic("bush_pref", theta[:, 0] * n / m)
    
    # to calculate probability that support for Bush shifted from debate 1 [0] to 2 [1]
    bush_shift = pm.Deterministic("bush_shift", bush_pref[1]-bush_pref[0])
    
    # because of the shapes of the inputs, this essentially creates 2 multinomials, 
    # one for each debate
    responses = pm.Multinomial("responses", n=n, p=theta, observed=y)
Explanation: Thus, we need to initialize a Dirichlet distribution prior with shape (2,3) and then refer to the relevant parameters by index where needed.
End of explanation
for v in polling_model_debates.unobserved_RVs:
    print(v, v.tag.test_value.shape)
Explanation: For models with multi-dimensional shapes, it's always good to check the shapes of the various parameters before sampling:
End of explanation
pm.model_to_graphviz(polling_model_debates)
Explanation: The plate notation visual can also help with that:
End of explanation
with polling_model_debates:
    polling_trace_debates = pm.sample(draws=3000, tune=1500)
with polling_model_debates:
    pm.traceplot(polling_trace_debates, combined=True)
Explanation: Let's sample with a slightly higher number of draws and tuning steps:
End of explanation
s = ["pre", "post"]
candidates = data["candidate"].values
pd.DataFrame(polling_trace_debates["theta"].mean(axis=0), index=s, columns=candidates)
Explanation: We'll take a look at the means of the posteriors for theta, indicating the % of support for each candidate pre & post debate:
End of explanation
pd.DataFrame(polling_trace_debates["bush_pref"].mean(axis=0), index=s, columns=["bush_pref"])
Explanation: Just from the means, we can see that the number of Bush supporters has likely decreased post debate from 48.8% to 46.3% (as a % of supporters of the 2 major candidates):
End of explanation
_, ax = plt.subplots(2,1, figsize=(10, 10))
sns.distplot(polling_trace_debates["bush_pref"][:,0], hist=False, ax=ax[0], label="Pre-Debate")
sns.distplot(polling_trace_debates["bush_pref"][:,1], hist=False, ax=ax[0], label="Post-Debate")
ax[0].set_title("% Responses for Bush vs Dukakis")
ax[0].set_xlabel("% Responses");
sns.distplot(polling_trace_debates["bush_shift"], hist=True, ax=ax[1], label="P(Bush Shift)")
ax[1].axvline(0, c='g', linestyle='dotted')
ax[1].set_title("% Shift Pre/Prior Debate")
ax[1].set_xlabel("% Shift");
Explanation: Let's compare the results visually, by plotting the posterior distributions of the pre/post debate values for % responses for Bush and the posterior for pre/post difference in Bush supporters:
End of explanation
perc_shift = (len(polling_trace_debates["bush_shift"][polling_trace_debates["bush_shift"] > 0])
              /len(polling_trace_debates["bush_shift"])
             )
print(f'P(Shift Towards Bush) = {perc_shift:.1%}')
Explanation: From the second plot, we can already see that a large portion of the posterior density is below 0, but let's be precise and actually calculate the probability that support shifted towards Bush after the debate:
End of explanation
<END_TASK>