---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.10.3
kernelspec:
  display_name: Python 3
  language: python
  name: pymc
---

```{code-cell} ipython3
import pandas as pd
import statsmodels as sm
import statsmodels.formula.api as smf
import pymc as pm
import numpy as np
```

# Ordinary least squares using PyMC

Recall the ordinary least squares model you might be familiar with
from a previous statistics or econometrics class.  We are interested
in estimating the parameter vector $\beta$ from the following
population regression function:

$$
\mathbf{y = x \beta + \epsilon}
$$

where $\mathbf{y}$ and $\mathbf{\epsilon}$ are each $N \times 1$
vectors of independent (outcome) and unobservable random variables.
The $N \times K$ matrix $\mathbf{x}$ are the non-random (assumed for
this notebook) independent variables (sometimes referred to as
covariates).  The least squares approach yields the "ols" estimator

$$
\mathbf{b} = \mathbf{{x'x}^{-1} x'y}
$$

The frequentist approach performs inference based on the variance
covariance matrix of $\mathbf{b}$ which is estimated as 

$$
cov(\mathbf{b}) = \hat{\sigma}^2 \left( \mathbf{x'x} \right)^{-1}
$$

where $\hat{\sigma}^2 = \frac{(\mathbf{y-xb})'(\mathbf{y-xb})}{N-K}$

The square root of the diagonal of this matrix yields standard errors
for estimated coefficients as reported in programs like Stata.  These
standard errors are used for inference related questions.

The Bayesian approach to OLS takes a different approach.  Consider the model

$$
\begin{align}
\sigma &\sim Uniform(a=.01, b=10) \nonumber \\
\mathbf{\beta}_k &\sim N(\mu_b = 0, \sigma_b = 10) \text{ } \forall k \in K \nonumber\\ 
\mathbf{y} &\sim N(\mathbf{x\beta}, \sigma^2 I) \nonumber 
\end{align}
$$

We can setup the likelihood and priors for getting information about
the posterior.  Assembling the pieces,

1. Likelihood:

   $$
   prob(\mathbf{y}| \mathbf{x}, \mathbf{\beta}, \mathbf{\sigma}) = \prod_{i=1}^N \frac{1}{\sqrt{2\pi\sigma^2}} e^\frac{(y_i-\mathbf{x}_i \mathbf{b})^2}{2\sigma^2} 
   $$
   
2. Prior on $\sigma$:
 
   $$
   prob(\sigma | a=0.01, b=10) = \frac{1}{(10 - .01)} \text{ for } .01  \le \sigma \le 10 \text{, 0 otherwise}
   $$
3. Prior on $\mathbf{\beta}$

   $$
   prob(\mathbf{b} | \mu_b = 0, \sigma_b = 10) = \prod_{k = 1}^{K} \frac{1}{\sqrt{2\pi \sigma^2_b}} e ^ {\frac{b_k - \mu_b}{2\sigma^2_b}} 
   $$


Let $\theta$ be the vector of all unknown parameters, $\begin{bmatrix} \beta_1 \\ \vdots \\ \beta_k \\ \vdots \\ \beta_K \\ \sigma  \end{bmatrix}$.  Then the posterior is proportional to 

$$
p(\theta | \mathbf{y}, \mathbf{x}, a, b, \mu_b, \sigma_b) \propto prob(\mathbf{y}| \mathbf{x}, \mathbf{\beta}, \mathbf{\sigma}) \times prob(\sigma | \alpha_l, \alpha_u) \times prob(\mathbf{b} | \mu_b = 0, \sigma_b = 10)
$$

## An example 

We will use the Tobias and Koop dataset for this excercise.  We'll
estimate a simple model using frequentist and Bayesian methods.  This
dataset has the following fields:

| Variable    | Definition                                    |
|-------------|-----------------------------------------------|
| id          | person id, ranging from 1 to 2,178            |
| educ        | education (years)                             |
| ln_wage     | log of hourly wage                            |
| pexp        | potential experience                          |
| time        | time trend                                    |
| ability     | ability                                       |
| meduc       | mother's education (years)                    |
| feduc       | father's education                            |
| broken_home | dummy variable for residence in a broken home |
| siblings    | number of siblings                            |
| pexp2       | squared potential experience                  |


```{code-cell} ipython3
tobias_koop = pd.read_stata('https://rlhick.people.wm.edu/econ407/data/tobias_koop.dta')
tobias_koop.head()
```

Note, that this is a panel dataset having the following time periods:

```{code-cell} ipython3
tobias_koop.time.value_counts()
```

For purposes of this lecture we will ignore the panel aspects of this
dataset and estimate an ordinary least squares model.  Furthermore,
due to memory constraints that students may encounter on `jupyterhub`
for some of the features below, we will restrict our sample to
`time==4` (1983). In general we don't recommend omitting data like
this, but do it here for the aforementioned reasons.  Reducing the
data:

```{code-cell} ipython3
print("Records in original dataset: ", tobias_koop.shape[0])
tobias_koop = tobias_koop[tobias_koop.time == 4]
print("Records in reduced dataset: ", tobias_koop.shape[0])
```

Let's estimate a simple wage model:

$$
ln\_wage_i = \beta_0 + \beta_1 educ + \beta_2 pexp + \beta_3 broken\_home_i + \epsilon_i
$$

Using statsmodels, we have

```{code-cell} ipython3
res = smf.ols(formula='ln_wage ~ educ + pexp + broken_home', 
              data=tobias_koop).fit()
print(res.summary())
```

### Bayesian Estimation

We will continue using the model setup as specified above,

$$
\begin{align}
\sigma &\sim Uniform(a=.01, b=10) \nonumber \\
\mathbf{\beta} &\sim N(\mu_b = 0, \sigma_b = 10) \nonumber \\
\mathbf{y} &\sim N(\mathbf{x\beta}, \sigma^2 I) \nonumber 
\end{align}
$$

The implementation in PyMC in code is very close to this model definition:

```{code-cell} ipython3
with pm.Model() as ols_model:
    
    # priors
    b0 = pm.Normal('constant', mu=0., sigma=10.)
    b1 = pm.Normal('educ', mu=0., sigma=10.)
    b2 = pm.Normal('pexp', mu=0., sigma=10.)
    b3 = pm.Normal('broken_home', mu = 0., sigma=10.)
    sigma = pm.Uniform('sigma', lower=.1, upper=10.)

    # yhat (or the mean for each i's likelihood contrib)
    mu_i = b0 + b1 * tobias_koop.educ + b2 * tobias_koop.pexp +\
                b3 * tobias_koop.broken_home
        
    # likelihood
    like = pm.Normal('likelihood', mu=mu_i, sigma=sigma, 
                      observed = tobias_koop.ln_wage)
```

Note, that MAP (maximum of the posterior) is very close to the OLS estimate:

```{code-cell} ipython3
with ols_model:
    startvals = pm.find_MAP()
    
print(startvals)
```


```{note}
On our jupyterhub server, it is necessary to include two function arguments, which in normal cases you might not want to use:

1.  Since each user is allocated 2 CPU cores.  For PyMC to run properly, you must use the `cores=2` argument below.  While the code will run without this argument, results may be unreliable particularly for this notebook.  On a typical PC, you would want to omit the `cores` argument and let PyMC use the maximum number of cores available for quickest execution.
2. The `progressbar=False` is included so that these notes will look better on this webpage.  In general, you would probably want to view the progress and wouldn't need this option.
```

Sampling from the posterior for the model is accomplished using

```{code-cell} ipython3
with ols_model:
    step = pm.Metropolis()
    samples = pm.sample(1000, step=step, tune=500, initvals=startvals, cores=2, chains=4,
                       return_inferencedata=True, progressbar=False)
```

Summarizing the trace:
```{code-cell} ipython3
pm.summary(samples)
```

The summary above shows

1. Results very similar to OLS
2. That we have very doubtful chain convergence as the Gelman-Rubin score (`r_hat` is large)

Examining autocorrelation, we have

```{code-cell} ipython3
pm.plot_autocorr(samples, combined=True);
```

which shows very high level of autocorrelation particularly for `educ`
and the model constant.

The accept rate is
```{code-cell} ipython3
accept = samples.sample_stats.accepted.mean().to_numpy()
print("Acceptance Rate: ", accept.round(3))
```

To obtain better results, we need to either

1. Increase sample size (and tuning), or
2. Choose a different step method

At this point in the course, our only option is to increase sampling
and tuning:

```{code-cell} ipython3
with ols_model:
    step = pm.Metropolis()
    samples = pm.sample(10000, step=step, tune=10000, initvals=startvals, cores=2, chains=4,
                       return_inferencedata=True, progressbar=False)
```

We get better, albeit not acceptable, convergence diagnostics:

```{code-cell} ipython3
pm.summary(samples)
```

but we still have high levels of autocorrelation:

```{code-cell} ipython3
pm.plot_autocorr(samples, combined=True);
```

Our acceptance rate is also improved:

```{code-cell} ipython3
accept = samples.sample_stats.accepted.mean().to_numpy()
print("Acceptance Rate: ", accept.round(3))
```

Here is the trace plot:

```{code-cell} ipython3
pm.plot_trace(samples);
```

### Adding "policy relevant" measures to your model

It is possible to add policy relevant measures in your model **and**
sample from those posteriors quite easily.  For demonstration
purposes, suppose we want to sample (1) the elasticity of education,
and (2) the predicted wage.  This is fairly straightforward but we
need to redefine our PyMC model to include these items:

Recall that given the $t^{th}$ sample for $\beta_{kt}$ in our trace the
elasticity in a semi-log model such as ours for individual $i$ and
variable $k$ is

$$
E_{ikt} = b_{kt} x_{ik}
$$

The average elasticity for variable $k$ and and sample $t$ is

$$
E_{kt} = \frac{\sum_{i=1}^{N} E_{ikt}}{N}
$$


The expected elasticity of variable $k$ as reflected by the underlying
uncertainty in our posterior is

$$
E_{k} = \int E_{kt} p(\theta | x_i, h^0) d\theta 
$$ 

where $\theta$ is the parameter vector and $p(\theta|x_i, h^0)$ is our
posterior and $h^0$ are hyperparameters describing the priors.

We can approximate this integral using samples from the posterior.
Since we have numerous samples for $b_k$, we estimate the integral above
using

$$
\hat{E}_k = \frac{\sum_{t=1}^T E_{kt}}{T}
$$

where $T$ is the number of samples in our trace.

To implement this in code, we will ask PyMC to store these
calculations in our sampled chain which we can average to produce the
desired measure.  Note the use of `pytensor.tensor` for performing these
calculations and the syntax `pm.Deterministic` tells `pymc` to
include the calculated values in the trace.

```{code-cell} ipython3
import pytensor.tensor as tt

with pm.Model() as ols_model:
    ## priors
    b0 = pm.Normal('constant', mu=0., sigma=10.)
    b1 = pm.Normal('educ', mu=0., sigma=10.)
    b2 = pm.Normal('pexp', mu=0., sigma=10.)
    b3 = pm.Normal('broken_home', mu=0., sigma=10.)
    sigma = pm.Uniform('sigma', lower=.1, upper=10.)
 
    ## yhat (or the mean for each i's likelihood contrib)
    mu_i = b0 + b1 * tobias_koop.educ + b2 * tobias_koop.pexp +\
                b3 * tobias_koop.broken_home
        
    ## include some policy relevant measures in our samples
    # for predicted wage, note the model is in log wage units
    yhat = pm.Deterministic('wage_hat', tt.exp(mu_i))
    # this is E_{ikt} for a semi log model for each
    # sample from our posterior (N x T):
    elas_educ = b1 * tobias_koop.educ
    # note this will be calculated for **each** observation
    # for **each** sampled parameter vector.  What we want is
    # the average over all observations for each sampled parameter
    # vector (E_{kt}) whick we can calculate as 
    elas = pm.Deterministic('elas_educ', tt.mean(elas_educ, axis=-1))
 
    # likelihood
    like = pm.Normal('likelihood', mu=mu_i, sigma=sigma, 
                      observed=tobias_koop.ln_wage)
```

Now let's sample from our posterior with these additional elements added in:

```{code-cell} ipython3
with ols_model:
    step = pm.Metropolis()
    samples = pm.sample(10000, step=step, tune=10000, initvals=startvals, cores=2, chains=4,
                       return_inferencedata=True, progressbar=False)
```

```{note}
On some computers, this may not run due to memory limitations or could take a long time to run.  If it fails to run on your machine, lower `10000` to some smaller value.
```

If we examine the `samples` object, note we have additional items:

```{code-cell} ipython3
samples.posterior
```

In particular, note that we have a new dimension (or coordinate) in
our chain, `wage_hat_dim_0` which tracks observations in our data.
Note, predicted wage is a large matrix:

```{code-cell} ipython3
samples.posterior.wage_hat.to_numpy().shape
```

Since we only have one elasticity measure per sampled parameter vector, it is 

```{code-cell} ipython3
samples.posterior.elas_educ.to_numpy().shape
```

in a similar manner to our other parameters:

```{code-cell} ipython3
samples.posterior.constant.to_numpy().shape
```

In our calculations, we elected **not** to include individual-specific
elasticities in our samples, hence there is no dimension over the
observations in our data.

### Analyzing the Trace

Summarizing the trace (and omitting our predicted wage) yields:

```{code-cell} ipython3
pm.summary(samples, var_names = ['constant', 'educ', 'pexp', 'broken_home', 'sigma', 'elas_educ'])
```

We can even see the marginal posterior for our elasticity estimate:

```{code-cell} ipython3
pm.plot_trace(samples, var_names = ['elas_educ']);
```

The predicted wage is included in the trace and is of dimension number
of chains x number of samples x number of individuals.  This is alot
of information.  To illustrate the information we have stored in our
posterior let's focus on only the first 50 individuals:

```{code-cell} ipython3
individs = [i for i in range(50)]
pm.plot_trace(samples.sel(wage_hat_dim_0 = individs), var_names=['wage_hat']);
```

By storing these values in our trace, we can calculate policy relevant
measures that **immediately** reflect the uncertainty of our parameter
estimates since the distribution is being driven by the distribution
of the posterior itself.  This is evident in the posterior
distributions of predicted wage for each individual.

## Summary

This workbook has demonstrated how to use `pymc` for a linear
regression model.  We have examined some convergence criteria and see
that we have to increase sample size quite alot to meet the
Gelman-Rubin metric of `r_hat` $\approx 1$.  Even after increasing the
number of samples by quite alot, the values above fail to completely
satisfy this criteria for some of the parameters in our model.

In the next chapter we turn our attention to more modern step methods
for achieving convergence with many fewer samples.