# Fitting a straight line to data

_Inspired by [Hogg et al. 2010](https://arxiv.org/abs/1008.4686) and [@jakevdp's notes](https://github.com/jakevdp/ESAC-stats-2014)_.

Python imports we'll need later...

```python
import numpy as np
import matplotlib.pyplot as plt

%matplotlib inline
rnd = np.random.RandomState(seed=42)
```

---

# Intro and choice of objective function

I want to start with a problem that everyone is probably familiar with or has at least seen before. The problem is this: we observe $N$ independent data points $\boldsymbol{y}=\{y_1,y_2,...y_N\}$ with uncertainties $\boldsymbol{\sigma}=\{\sigma_1,\sigma_2,...\sigma_N\}$ at perfectly-measured values $\boldsymbol{x}=\{x_1,x_2,...x_N\}$. We have reason to believe that the these data were generated by a process that is well-represented by a straight-line, and the only reason that the data deviate from this straight line is because of uncorrelated, Gaussian measurement noise in the $y$-direction. Let's first generate some data that meet these qualifications:

```python
n_data = 16 # number of data points
a_true = 1.255 # randomly chosen truth
b_true = 4.507
```

---

### Exercise 1:

1. Randomly generate an array of uniformly-distributed `x` values from the domain `(0,2)`.
2. Sort the values in ascending order.

```python
# Fill in your solution here
```

Execute the code below and verify that it executes:

```python
# evaluate the true model at the given x values
y = a_true*x + b_true

# Heteroscedastic Gaussian uncertainties only in y direction
y_err = rnd.uniform(0.1, 0.2, size=n_data) # randomly generate uncertainty for each datum
y = y + rnd.normal(0, y_err) # add noise to y data
```

```python
plt.errorbar(x, y, y_err, marker='o', linestyle='none')
plt.xlabel('$x$')
plt.ylabel('$y$')
plt.tight_layout()
```

---

Now let's forget that we did that -- we know nothing about the model parameters, except that we think the true values of the data are well-described by a linear relation! We would like to measure the "best-fit" parameters of this model (for a straight line, the slope and intercept $(a,b)$) given the data above. In math, our model for the data $y$ is:
$$
\begin{align}
y &= f(x \,;\, a, b) + {\rm noise}\\
f(x \,;\, a, b) &= a\,x + b
\end{align}
$$

For a given set of parameters, $(a,b)$, we can evaluate our model $f(x \,;\, a, b)$ at a given $x$ location to compute the value of $y$ that we would expect in the absence of noise. For example, for the $n$th datum and for a given set of parameter values $(a,b)$:

$$
\tilde{y}_n = f(x_n \,;\, a, b)
$$

Now, we somehow want to search through all possible values of $a,b$ to find the "best" values, given the data, with some definition of "best." When we say this word, we are implying that we want to _optimize_ (find the maximum or minimum) some _objective function_ (a function that takes our data, our model, and returns a quantification of "best", usually as a scalar). Numerically, this scalar objective function can be any function (though you probably want it to be convex) and you will see different choices in practice. You have some leeway in this choice depending on whether your goal is _prediction_, _discovery_, or _data compression_.

However, for _inference_—the typical use-case for us as scientists—you don't have this freedom: one of the conclusions of this talk is going to be that __you have no choice about what "best" means__! Before we get there, though, let's explore what seem like reasonable choices.

Here are a few desirable features we'd like any objective function to have:

1. For a given set of parameters, we should compare our predicted values to the measured values and base our objective function on the differences
2. The scalar value should be dimensionless (the value of the objective function shouldn't care if we use kilometers vs. parsecs)
3. Data points that have larger errors should contribute less to the objective function (if a datum has a large offset from the predicted value, it shouldn't matter _if_ the datum has a large uncertainty)
4. Convexity

To meet these three criteria, whatever objective function we choose should operate on the (dimensionless) quantities:

$$
\chi_n = \frac{y_n - \tilde{y}_n(x_n; a,b)}{\sigma_n}
$$

i.e. the difference between our predicted values $\tilde{y}$ and the observed $y$ values, weighted by the inverse uncertainties $\sigma$. The uncertainties have the same units as the data, so this is a dimensionless quantity. It also has the nice property that, as we wanted, points with large uncertainties are _downweighted_ relative to points with small uncertainties. Here are some ideas for objective functions based on this scalar:

- __Weighted absolute deviation__: the sum of the absolute values

    $\sum_n^N \, \left|\chi_n\right|$

- __Weighted squared deviation__: the sum of the squares

    $\sum_n^N \, \chi_n^2$

- __Weighted absolute deviation to some power__ $p$:

    $\sum_n^N \, \left|\chi_n\right|^p $

_(Note: don't show this to statisticians or they will get me fired. To a statistician, $\chi^2$ is a distribution not a statistic...but astronomers seem to use this terminology.)_

For simplicity, let's just compare two of these: the absolute deviation and the squared deviation.

---

### Exercise 2:

Implement the functions to compute the weighted deviations below

```python
# FILL IN THESE FUNCTIONS:

def line_model(pars, x):
    pass

def weighted_absolute_deviation(pars, x, y, y_err):
    pass

def weighted_squared_deviation(pars, x, y, y_err):
    pass
```

Verify that you've correctly implemented your functions by executing the following cell:

```python
_pars = [1., -10.]
_x = np.arange(16)
_y = _x
_yerr = np.ones_like(_x)

truth = np.array([-10.,  -9.,  -8.,  -7.,  -6.,  -5.,  -4.,  -3.,  -2.,  -1.,   0., 1.,   2.,   3.,   4.,   5.])
assert np.allclose(line_model(_pars, _x), truth), 'Error in line_model() function!'
assert weighted_absolute_deviation(_pars, _x, _y, _yerr) == 160., 'Error in weighted_absolute_deviation() function!'
assert weighted_squared_deviation(_pars, _x, _y, _yerr) == 1600., 'Error in weighted_squared_deviation() function!'
```

---

We can demonstrate that these are convex (over some domain) by computing the objective function values over a grid of parameter values (a grid in $a, b$):

```python
# make a 256x256 grid of parameter values centered on the true values
a_grid = np.linspace(a_true-2., a_true+2, 256)
b_grid = np.linspace(b_true-2., b_true+2, 256)
a_grid,b_grid = np.meshgrid(a_grid, b_grid)
ab_grid = np.vstack((a_grid.ravel(), b_grid.ravel())).T
```

```python
fig,axes = plt.subplots(1, 2, figsize=(9,5.1), sharex=True, sharey=True)

for i,func in enumerate([weighted_absolute_deviation, weighted_squared_deviation]):
    func_vals = np.zeros(ab_grid.shape[0])
    for j,pars in enumerate(ab_grid):
        func_vals[j] = func(pars, x, y, y_err)

    axes[i].pcolormesh(a_grid, b_grid, func_vals.reshape(a_grid.shape),
                       cmap='Blues', vmin=func_vals.min(), vmax=func_vals.min()+256) # arbitrary scale

    axes[i].set_xlabel('$a$')

    # plot the truth
    axes[i].plot(a_true, b_true, marker='o', zorder=10, color='#de2d26')
    axes[i].axis('tight')
    axes[i].set_title(func.__name__, fontsize=14)

axes[0].set_ylabel('$b$')

fig.tight_layout()
```

There are minima in both cases near the true values of the parameters (good), but the functions clearly look different. Which one should we choose for finding the best parameters?

In order to pick between these two, or any of the arbitrary objective functions we could have chosen, we have to _justify_ using one function over the others. In what follows, we'll justify optimizing the sum of the squared deviations (so-called "least-squares fitting") by thinking about the problem _probabilistically_, rather than procedurally.

### Least-squares fitting

Let's review the assumptions we made above in generating our data:

1. The data were generated by a straight line
2. Uncorrelated, _known_ Gaussian uncertainties in $y$ cause deviations between the data and predictions
3. The data points are independent
4. The $x$ data are known perfectly, or at least their uncertainties are _far smaller_ than the uncertainties in $y$

First off, these assumptions tell us that for each datum $(x_n, y_n)$ there is some true $y_{n,{\rm true}}$, and because of limitations in our observing process we can't observe the truth, but we know that the values we do observe will be Gaussian (Normal) distributed around the true value. _(Note: This assumption tends to be a good or at least a conservative approximation in practice, but there are certainly more complex situations when, e.g., you have asymmetric uncertainties, or error distributions with large tails!)_. In math:

$$
\begin{align}
p(y \,|\, y_{\rm true}) &= \mathcal{N}(y \,|\, y_{\rm true}, \sigma^2) \\
\mathcal{N}(y \,|\, y_{\rm true}, \sigma^2) &= (2\pi \sigma^2)^{-1/2} \, \exp\left(-\frac{1}{2} \frac{(y-y_{\rm true})^2}{\sigma^2} \right)
\end{align}
$$

This is the likelihood of observing a particular $y$ given the true $y_{\rm true}$. Note that in our model, all of the $y_{\rm true}$'s must lie on a line. It is also interesting that the argument of the normal distribution looks a lot like $\chi^2$!

What about considering two data points, $y_1$ and $y_2$? Now we need to write down the _joint_ probability

$$
p(y_1, y_2 \,|\, y_{1,{\rm true}}, \sigma_1, y_{2,{\rm true}}, \sigma_2)
$$

But, note that in assumption 3 above, we are assuming the data are independent. In that case, the random error in one point does not affect the random error in any other point, so the joint probability can be turned into a product:

$$
p(\{y_n\} \,|\, \{y_{n,{\rm true}}\}, \{\sigma_n\}) = \prod_n^N \, p(y_n \,|\, y_{n,{\rm true}}, \sigma_n)
$$

This is the full expression for the likelihood of the observed data given the true $y$ values. Recall that these true values, according to our assumptions, must lie on a line with some parameters, and we're trying to infer those parameters! We can compute a particular $y_{n,{\rm true}}$ using $x_n$ and a given set of model parameters $a, b$. With that in mind, we can write the likelihood instead as:

$$
p(\{y_n\} \,|\, a, b, \{x_n\}, \{\sigma_n\}) = \prod_n^N \, p(y_n \,|\, a, b, x_n, \sigma_n)
$$

So what are the "best" values of the parameters $a, b$? They are the ones that _maximize_ this likelihood!

The product on the right of the likelihood is a product over exponentials (well, Gaussians), which can be annoying to deal with. But, maximizing the likelihood is equivalent to maximizing the _log_-likelihood -- so we can get rid of the product and all of those exponentials by taking the log of both sides:

$$
\begin{align}
\ln p(\{y_n\} \,|\, a, b, \{x_n\}, \{\sigma_n\}) &= \sum_n^N \, \ln\left[p(y_n \,|\, a, b, x_n, \sigma_n)\right] \\
&= \sum_n^N \ln \left[(2\pi \sigma_n^2)^{-1/2} \,
    \exp\left(-\frac{1}{2} \frac{(y_n-(a\,x_n+b))^2}{\sigma_n^2} \right) \right] \\
&= -\frac{N}{2}\ln(2\pi)
    - \frac{1}{2} \sum_n^N \left[\frac{(y_n-(a\,x_n+b))^2}{\sigma_n^2} + \ln{\sigma_n^2} \right]
\end{align}
$$

In this case, the uncertainties are known and constant, so to maximize this expression we only care that (abbreviating the likelihood as $\mathcal{L}$):

$$
\begin{align}
\ln \mathcal{L} &= - \frac{1}{2} \sum_n^N \left[\frac{(y_n-(a\,x_n+b))^2}{\sigma_n^2}\right] + {\rm const.} \\
&= - \frac{1}{2} \sum_n^N \, \chi_n^2 + {\rm const.} \\
\end{align}
$$

Apparently, _minimizing_ the sum of the weighted squared deviations is equivalent to _maximizing_ the (log) likelihood derived from thinking about the probability of the data! That is great because (a) it directly gives us the uncertainties on the inferred model parameters, and (b) it's an analytic way to solve this problem using linear algebra which is _really_ fast!

### Least-squares / maximum likelihood with matrix calculus

Using linear algebra, we can simplify and generalize a lot of the expressions above. In what follows, all vectors are column vectors and are represented by lower-case bold symbols. Matrices are upper-case bold symbols.

We'll start by writing our model as a matrix equation. To do that, we need a way to, for a given set of parameters, compute the set of predicted $y$'s. This is done by defining the parameter vector, $\boldsymbol{\theta}$, and a matrix typically called the _design matrix_, $\boldsymbol{X}$:

$$
\boldsymbol{\theta} = \begin{bmatrix} b \\ a \end{bmatrix} \quad
\boldsymbol{X} = \begin{bmatrix} 1 & x_1 \\ 1 & x_2 \\ \vdots & \vdots \\ 1 & x_N \end{bmatrix}
$$

(note the order of the parameters!). With these definitions, the vector of predicted $y$ values is just

$$
\boldsymbol{y}_{\rm pred} = \boldsymbol{X} \, \boldsymbol{\theta}
$$

so the deviation vector between the prediction and the data is just $(\boldsymbol{y}-\boldsymbol{X} \, \boldsymbol{\theta})$ where

$$
\boldsymbol{y} = \begin{bmatrix} y_1 \\ y_2 \\ \vdots \\ y_N \end{bmatrix}
$$

But how do we include the uncertainties? We'll pack the list of uncertainties (variances) into the trace of a 2D, $N \times N$ matrix called the _covariance matrix_. Because we are assuming the uncertainties are independent, the off-diagonal terms are all zero:

$$
\boldsymbol{\Sigma} = \begin{bmatrix}
\sigma_1^2 & 0 & \dots & 0 \\
0 & \sigma_2^2 & \dots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & 0 & \sigma_N^2
\end{bmatrix}
$$

With these matrices, we can write the expression for $\chi^2$ (and therefore the log-likelihood) very concisely:

$$
\begin{align}
\chi^2 &= \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right)^\mathsf{T} \,
    \boldsymbol{\Sigma}^{-1} \,
    \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right) \\
\ln\mathcal{L} &= -\frac{1}{2}\left[N\,\ln(2\pi)
    + \ln|\boldsymbol{\Sigma}|
    + \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right)^\mathsf{T} \,
      \boldsymbol{\Sigma}^{-1} \,
      \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right)
\right]
\end{align}
$$

In this form, the terms in the $\chi^2$ have a nice geometric interpretation: This looks like a distance between the data and the model computed with the metric $\boldsymbol{\Sigma}$.

If you solve for the optimum of the log-likelihood function (take the derivative with respect to $\boldsymbol{\theta}$ and set equal to 0), you find that:

$$
\newcommand{\trpo}[1]{{#1}^{\mathsf{T}}}
\newcommand{\bs}[1]{\boldsymbol{#1}}
\bs{\theta}_{\rm best} = \left[\trpo{\bs{X}} \, \bs{\Sigma}^{-1} \, \bs{X}\right]^{-1} \,
    \trpo{\bs{X}} \, \bs{\Sigma}^{-1} \, \bs{y}
$$

Getting the best-fit parameters just requires a few simple linear algebra operations! As an added bonus, we also get the _uncertainties_ on the parameters. The $2\times2$ covariance matrix for the best-fit parameters is given by the matrix:

$$
\newcommand{\trpo}[1]{{#1}^{\mathsf{T}}}
\newcommand{\bs}[1]{\boldsymbol{#1}}
C = \left[\trpo{\bs{X}} \, \bs{\Sigma}^{-1} \, \bs{X}\right]^{-1}
$$

That means we can just write out the linear algebra explicitly and use `numpy.linalg` to solve it for us!

### Exercise 3:

Implement the necessary linear algebra to solve for the best-fit parameters and the parameter covariance matrix, defined above. Call these `best_pars` and `pars_Cov`, respectively.

```python

```

Now let's look at the covariance matrix of the parameters (the uncertainty in the parameters) and plot the 1 and 2-sigma error ellipses:

```python
# some tricks to get info we need to plot an ellipse, aligned with
#    the eigenvectors of the covariance matrix
eigval,eigvec = np.linalg.eig(pars_Cov)
angle = np.degrees(np.arctan2(eigvec[1,0], eigvec[0,0]))
w,h = 2*np.sqrt(eigval)
```

```python
from matplotlib.patches import Ellipse

fig,ax = plt.subplots(1, 1, figsize=(5,5))

for n in [1,2]:
    ax.add_patch(Ellipse(best_pars, width=n*w, height=n*h, angle=angle,
                         fill=False, linewidth=3-n, edgecolor='#555555',
                         label=r'{}$\sigma$'.format(n)))

ax.plot(b_true, a_true, marker='o', zorder=10, label='truth')
ax.plot(best_pars[0], best_pars[1], marker='o', zorder=9, label='estimate')

ax.set_xlabel('$b$')
ax.set_ylabel('$a$')
ax.legend(loc='best')

fig.tight_layout()
```

There we have it! The best-fit parameters and their errors for the straight-line fit, optimized with the only justifyable objective function, directly from a few linear algebra calculations.

This approach can be generalized somewhat (e.g. to account for correlated errors as off-diagonal elements in the covariance matrix $\Sigma$), but it only works **for models with linear parameters**, meaning that parameters enter linearly in the model function $f$.

---
# Bonus material (peruse at your leisure)

## The Bayesian approach

Let's review what we did so far. We found that standard weighted least squares fitting is a justified approach to estimating the best-fit parameters because it optimizes the likelihood of the data under the assumptions of our model; it optimizes a _justified scalar objective function_. We then fit our straight-line model to the data using and got back a point-estimate of the best parameters along with a covariance matrix describing the uncertainties in the parameters. This is the way of the _frequentist_. What we're going to do now is see what happens if we switch to a Bayesian methodology instead. While the two methods end up looking mathematically identical, there are fundamental philosophical differences that can lead to very different interpretations and implementations when models are more complex than the toy example we use above.

As Bayesians, we aren't interested in a point-estimate of the best parameters, but rather we're interested in the inferred distribution of possible parameter values (the _posterior probability distribution function_ over parameters). So how do we write down or solve for this posterior pdf? Before we get to that, let's take a look at a fundamental equation of Bayesian statistics, [Bayes' theorem](https://en.wikipedia.org/wiki/Bayes'_theorem), which we'll derive using the joint probability of $A$ and $B$ which are conditional on some other information $I$ that, right now, we don't care about. For example, $A$ could be the time it takes to get from here to NYC, $B$ could be the amount of traffic on the road, and $I$ could include the information that we're driving a car and not walking. Bayes' theorem as expressed below is not controversial -- Bayesians and Frequentists agree that this is just how joint and conditional probabilities work. We start by writing down the joint probability of $A$ and $B$, then factor it in two ways into conditional proabilities:

$$
p(A,B \,|\, I) = p(A\,|\,B, I)\,p(B\,|\, I) = p(B\,|\,A, I)\,p(A \,|\, I)
$$

Now we look at the right two expressions, and divide by one of the marginal probabilities to get:

$$
p(A\,|\,B, I) = \frac{p(B\,|\,A, I)\,p(A \,|\, I)}{p(B\,|\, I)}
$$

Ok, so that's all fine. Now let's replace $A$ and $B$ with variables that represent, from our example above, our data $D=(\{x_n\},\{y_n\},\{\sigma_n\})$ and our model parameters $\boldsymbol{\theta}$:

$$
p(\boldsymbol{\theta}\,|\,D, I) = \frac{p(D\,|\,\boldsymbol{\theta}, I)\,p(\boldsymbol{\theta} \,|\, I)}{p(D\,|\, I)}
$$

In just switching the meaning of the variables, this expression becomes controversial! Frequentists would object to the above for two main reasons:

1. The term on the left hand side is a probability over parameters given the data (the _posterior_ pdf) $p(\boldsymbol{\theta}\,|\,D, I)$. This is something that a frequentist would say cannot exist - there is only one true vector of parameters that we are trying to learn, not a distribution!
2. The right-most term in the numerator is a probability over parameters _with no dependence on the data_ (the _prior_ pdf). This encapsulates all of our prior knowledge about the parameters before we did the experiment and observed some data. This is perhaps the aspect of Bayesian inference that frequentists most disagree with.

The differences above result from the fact that probability means something different to Frequentists and Bayesians. Bayesians think of probability as representing a _degree of belief_ about something, whereas a frequentist thinks of a probability as related to _limiting frequencies of occurrence_ in repeated trials or observations. This is a rich topic and I highly recommend reading [this series of blogposts](http://jakevdp.github.io/blog/2014/03/11/frequentism-and-bayesianism-a-practical-intro/) by Jake Vanderplas to learn more. For now, let's put on Bayesian hats and take a look at the implications of the expression above.

(_It's good to rememeber that we're all friends. The differences are based on philosophy and so can lead to some heated discussions and debates, but remember we're all trying to do science -- we're on the same team!_)

## Bayes' theorem and Bayesian inference

Let's decompose Bayes' theorem (as applied to modeling and inference). The four terms in Bayes' theorem above have names that are good to be familiar with:

- $p(\boldsymbol{\theta}\,|\,D, I)$ - __posterior probability__:
    This is the thing we are after when we do Bayesian inference or model fitting. We want to know what the distribution of possible parameter values is, given the data we observe and any prior information or assumptions $I$.

- $p(D\,|\,\boldsymbol{\theta}, I)$ - __likelihood__:
    This is the likelihood of the data given a particular set of model parameters. We've already seen this object and used it above to find the best-fit model parameters by maximizing this function. In a Bayesian context, it can also be thought of as a distribution -- it's a distribution that generates new datasets given a model instance. For that reason, we typically refer to models that produce a likelihood as _generative models_ because they specify how to  generate new data sets that look like the one you observe. As we saw above when we wrote the likelihood function for a straight line model and data with Gaussian errors, the likelihood usually contains a component that can be interpreted as the _noise model_.

- $p(\boldsymbol{\theta} \,|\, I)$ - __prior probability__
    This contains any relevant information about our parameters that we know before observing the data. This can include physical constraints, previous measurements, or anything, really. This flexibility is what makes the prior a somewhat controversial object. In practice, the prior only really matters if it is much narrower than the likelihood function. If the prior is broad with respect to the likelihood, the information in the likelihood makes the prior almost irrelevant. However, there are several subtleties to choosing priors that need to be considered. As an example, one subtlety comes from the choice of coordinates for the model parameters: a prior that is broad and flat in a parameter $\alpha$ won't be broad and flat if you change variables to $\beta = \alpha^2$.

- $p(D\,|\, I)$ - __evidence__ or __fully marginalized likelihood__ (FML)
    In many cases the evidence is simply a normalization constant and, for some of the most relevant algorithms used in inference, can be ignored. This term involves an integral over all of parameter space that can be very difficult to compute:

    $$
    p(D\,|\, I) = \int \,\mathrm{d}\boldsymbol{\theta} \, p(D\,|\,\boldsymbol{\theta}, I) \, p(\boldsymbol{\theta} \,|\, I)
    $$

    If you need to do Bayesian model selection (e.g., decide between models with different parameters), you unfortunately need to compute this quantity. But if you only _think_ you need the FML, beware!

So how do we make use of all of this, in practice?

Let's return to our example of fitting a line to data with the same data as above. In some sense, we are almost done once we write down an expression for the posterior pdf. If we ignore the FML, this amounts to multiplying a likelihood by a prior pdf. Well, we've already done the most important part: we already wrote down the likelihood function! This is often the hardest part and what we spend the most time doing as scientists (well, assuming you're not building the instrument to observe the data!). We now need to define a prior pdf over the model parameters. Here we have some flexibility. Two possibilities you can always consider:

1. A completely uninformative prior, based on dimensionality, symmetry, or entropy arguments (sometimes, this will mean using a _flat prior_ or _uniform prior_)
2. An empirical prior, based on previous _independent data_ that constrains this model (e.g., a previous measurement of the model parameters from an earlier dataset)

For simplicity, we're going to assume a flat prior over both slope and intercept. Note that for this problem, this is [_not_ an uninformative prior](http://jakevdp.github.io/blog/2014/06/14/frequentism-and-bayesianism-4-bayesian-in-python/). For now, we'll assume that the data are informative enough that the small bias we introduce by using this prior is negligible. Let's now define the functions we'll need, and recall that

$$
\ln\mathcal{L} = -\frac{1}{2}\left[N\,\ln(2\pi)
    + \ln|\boldsymbol{\Sigma}|
    + \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right)^\mathsf{T} \,
      \boldsymbol{\Sigma}^{-1} \,
      \left(\boldsymbol{y} - \boldsymbol{X}\,\boldsymbol{\theta}\right)
\right]
$$

### Exercise 4:

Implement the log-prior method (`ln_prior`) on the model class below.

#### Solution:

```python
class StraightLineModel(object):

    def __init__(self, x, y, y_err):
        """
        We store the data as attributes of the object so we don't have to
        keep passing it in to the methods that compute the probabilities.
        """
        self.x = np.asarray(x)
        self.y = np.asarray(y)
        self.y_err = np.asarray(y_err)

    def ln_likelihood(self, pars):
        """
        We don't need to pass in the data because we can access it from the
        attributes. This is basically the same as the weighted squared
        deviation function, but includes the constant normalizations for the
        Gaussian likelihood.
        """
        N = len(self.y)
        dy = self.y - line_model(pars, self.x)
        ivar = 1 / self.y_err**2 # inverse-variance
        return -0.5 * (N*np.log(2*np.pi) + np.sum(2*np.log(self.y_err)) + np.sum(dy**2 * ivar))

    def ln_prior(self, pars):
        """
        The prior only depends on the parameters, so we don't need to touch
        the data at all. We're going to implement a flat (uniform) prior
        over the ranges:
            a : [0, 100]
            b : [-50, 50]

        """

        a, b = pars # unpack parameters
        ln_prior_val = 0. # we'll add to this

        if a < 0 or a > 100.:
            return -np.inf
        else:
            ln_prior_val += np.log(1E-2) # normalization, log(1/100)

        if b < -50 or b > 50.:
            return -np.inf
        else:
            ln_prior_val += np.log(1E-2) # normalization, log(1/100)

        return ln_prior_val

    def ln_posterior(self, pars):
        """
        Up to a normalization constant, the log of the posterior pdf is just
        the sum of the log likelihood plus the log prior.
        """
        lnp = self.ln_prior(pars)
        if np.isinf(lnp): # short-circuit if the prior is infinite (don't bother computing likelihood)
            return lnp

        lnL = self.ln_likelihood(pars)
        lnprob = lnp + lnL

        if np.isnan(lnprob):
            return -np.inf

        return lnprob

    def __call__(self, pars):
        return self.ln_posterior(pars)
```

```python
model = StraightLineModel(x, y, y_err)
```

Now we'll repeat what we did above to map out the value of the log-posterior over a 2D grid of parameter values. Because we used a flat prior, you'll notice it looks identical to the visualization of the `weighted_squared_deviation` -- only the likelihood has any slope to it!

```python
def evaluate_on_grid(func, a_grid, b_grid, args=()):
    a_grid,b_grid = np.meshgrid(a_grid, b_grid)
    ab_grid = np.vstack((a_grid.ravel(), b_grid.ravel())).T

    func_vals = np.zeros(ab_grid.shape[0])
    for j,pars in enumerate(ab_grid):
        func_vals[j] = func(pars, *args)

    return func_vals.reshape(a_grid.shape)
```

```python
fig,axes = plt.subplots(1, 3, figsize=(14,5.1), sharex=True, sharey=True)

# make a 256x256 grid of parameter values centered on the true values
a_grid = np.linspace(a_true-5., a_true+5, 256)
b_grid = np.linspace(b_true-5., b_true+5, 256)

ln_prior_vals = evaluate_on_grid(model.ln_prior, a_grid, b_grid)
ln_like_vals = evaluate_on_grid(model.ln_likelihood, a_grid, b_grid)
ln_post_vals = evaluate_on_grid(model.ln_posterior, a_grid, b_grid)

for i,vals in enumerate([ln_prior_vals, ln_like_vals, ln_post_vals]):
    axes[i].pcolormesh(a_grid, b_grid, vals,
                       cmap='Blues', vmin=vals.max()-1024, vmax=vals.max()) # arbitrary scale

axes[0].set_title('log-prior', fontsize=20)
axes[1].set_title('log-likelihood', fontsize=20)
axes[2].set_title('log-posterior', fontsize=20)

for ax in axes:
    ax.set_xlabel('$a$')

    # plot the truth
    ax.plot(a_true, b_true, marker='o', zorder=10, color='#de2d26')
    ax.axis('tight')

axes[0].set_ylabel('$b$')

fig.tight_layout()
```

### Exercise 5:

Subclass the `StraightLineModel` class and implement a new prior. Replace the flat prior above with an uncorrelated 2D Gaussian centered on $(\mu_a,\mu_b) = (3., 5.5)$ with root-variances $(\sigma_a,\sigma_b) = (0.05, 0.05)$. Compare the 2D grid plot with the flat prior to the one with a Gaussian prior

#### Solution:

```python
class StraightLineModelGaussianPrior(StraightLineModel): # verbose names are a good thing!

    def ln_prior(self, pars):
        a, b = pars # unpack parameters
        ln_prior_val = 0. # we'll add to this

        # prior on a is a Gaussian with mean, stddev = (3, 0.05)
        ln_prior_val += -0.5*(a - 3.)**2/0.05**2 # this is not normalized properly, but that's ok

        # prior on b is a Gaussian with mean, stddev = (5.5, 0.05)
        ln_prior_val += -0.5*(b - 5.5)**2/0.05**2 # this is not normalized properly, but that's ok

        return ln_prior_val
```

```python
model_Gprior = StraightLineModelGaussianPrior(x, y, y_err)
```

```python
fig,axes = plt.subplots(1, 3, figsize=(14,5.1), sharex=True, sharey=True)

ln_prior_vals2 = evaluate_on_grid(model_Gprior.ln_prior, a_grid, b_grid)
ln_like_vals2 = evaluate_on_grid(model_Gprior.ln_likelihood, a_grid, b_grid)
ln_post_vals2 = evaluate_on_grid(model_Gprior.ln_posterior, a_grid, b_grid)

for i,vals in enumerate([ln_prior_vals2, ln_like_vals2, ln_post_vals2]):
    axes[i].pcolormesh(a_grid, b_grid, vals,
                       cmap='Blues', vmin=vals.max()-1024, vmax=vals.max()) # arbitrary scale

axes[0].set_title('log-prior', fontsize=20)
axes[1].set_title('log-likelihood', fontsize=20)
axes[2].set_title('log-posterior', fontsize=20)

for ax in axes:
    ax.set_xlabel('$a$')

    # plot the truth
    ax.plot(a_true, b_true, marker='o', zorder=10, color='#de2d26')
    ax.axis('tight')

axes[0].set_ylabel('$b$')

fig.tight_layout()
```

---

Now what do we do? The answer depends a bit on your intentions. If you'd like to propagate the posterior pdf (as in, pass on to other scientists to use your results), what do you do if the posterior pdf isn't analytic? And what numbers do you put in your abstract? One option is to draw samples from your posterior pdf and compute summary statistics (e.g., median and quantils) using the samples. That's the approach we're going to take.

## MCMC

One of the most common and powerful class of methods people use for generating these samples is Markov Chain Monte Carlo (MCMC), but there are other options (e.g., brute-force or monte carlo rejection sampling). MCMC methods are useful because they scale reasonably to higher dimensions (well, at least better than brute-force). A disadvantage to these methods comes from the "Markov Chain" part of the name: there is always some correlation between nearby steps in a chain of samples, so you have to compute second-order statistics on the samples to try to verify whether your samples are truly random or fair samples from the target distribution (your posterior pdf).

The simplest MCMC algorithm is known as Metropolis-Hastings. I'm not going to explain it in detail, but in pseudocode, it is:

- Start from some position in parameter space, $\theta_0$ with posterior probability $\pi_0$
- Iterate from 1 to $N_{\rm steps}$:
    - Sample an offset from $\delta\theta_0$ from some proposal distribution
    - Compute a new parameter value using this offset, $\theta_{\rm new} = \theta_0 + \delta\theta_0$
    - Evaluate the posterior probability at the new new parameter vector, $\pi_{\rm new}$
    - Sample a uniform random number, $r \sim \mathcal{U}(0,1)$
    - if $\pi_{\rm new}/\pi_0 > 1$ or $\pi_{\rm new}/\pi_0 > r$:
        - store $\theta_{\rm new}$
        - replace $\theta_0,\pi_0$ with $\theta_{\rm new},\pi_{\rm new}$
    - else:
        - store $\theta_0$ again

The proposal distribution has to be chosen and tuned by hand. We'll use a spherical / uncorrelated Gaussian distribution with root-variances set by hand:

```python
def sample_proposal(*sigmas):
    return np.random.normal(0., sigmas)

def run_metropolis_hastings(p0, n_steps, model, proposal_sigmas):
    """
    Run a Metropolis-Hastings MCMC sampler to generate samples from the input
    log-posterior function, starting from some initial parameter vector.

    Parameters
    ----------
    p0 : iterable
        Initial parameter vector.
    n_steps : int
        Number of steps to run the sampler for.
    model : StraightLineModel instance (or subclass)
        A callable object that takes a parameter vector and computes
        the log of the posterior pdf.
    proposal_sigmas : list, array
        A list of standard-deviations passed to the sample_proposal
        function. These are like step sizes in each of the parameters.
    """
    p0 = np.array(p0)
    if len(proposal_sigmas) != len(p0):
        raise ValueError("Proposal distribution should have same shape as parameter vector.")

    # the objects we'll fill and return:
    chain = np.zeros((n_steps, len(p0))) # parameter values at each step
    ln_probs = np.zeros(n_steps) # log-probability values at each step

    # we'll keep track of how many steps we accept to compute the acceptance fraction
    n_accept = 0

    # evaluate the log-posterior at the initial position and store starting position in chain
    ln_probs[0] = model(p0)
    chain[0] = p0

    # loop through the number of steps requested and run MCMC
    for i in range(1,n_steps):
        # proposed new parameters
        step = sample_proposal(*proposal_sigmas)
        new_p = chain[i-1] + step

        # compute log-posterior at new parameter values
        new_ln_prob = model(new_p)

        # log of the ratio of the new log-posterior to the previous log-posterior value
        ln_prob_ratio = new_ln_prob - ln_probs[i-1]

        if (ln_prob_ratio > 0) or (ln_prob_ratio > np.log(np.random.uniform())):
            chain[i] = new_p
            ln_probs[i] = new_ln_prob
            n_accept += 1

        else:
            chain[i] = chain[i-1]
            ln_probs[i] = ln_probs[i-1]

    acc_frac = n_accept / n_steps
    return chain, ln_probs, acc_frac
```

Now we'll run the sampler! Let's start from some arbitrary position allowed by our prior.

### Exercise 6:

Choose a starting position, values for `a` and `b` to start the MCMC from. In general, a good way to do this is to sample from the prior pdf. Generate values for `a` and `b` by sampling from a uniform distribution over the domain we defined above. Then, run the MCMC sampler from this initial position for 8192 steps. Play around with ("tune" as they say) the `proposal_sigmas` until you get an acceptance fraction around ~40%.

#### Solution:

```python
p0 = [6.,6.]
chain,_,acc_frac = run_metropolis_hastings(p0, n_steps=8192, model=model,
                                           proposal_sigmas=[0.05,0.05])
print("Acceptance fraction: {:.1%}".format(acc_frac))
```

Let's look at the chain returned, the parameter value positions throughout the sampler run:

```python
fig,ax = plt.subplots(1, 1, figsize=(5,5))

ax.pcolormesh(a_grid, b_grid, ln_post_vals,
              cmap='Blues', vmin=ln_post_vals.max()-128, vmax=ln_post_vals.max()) # arbitrary scale
ax.axis('tight')

fig.tight_layout()

ax.plot(a_true, b_true, marker='o', zorder=10, color='#de2d26')
ax.plot(chain[:512,0], chain[:512,1], marker='', color='k', linewidth=1.)

ax.set_xlabel('$a$')
ax.set_ylabel('$b$')
```

We can also look at the individual parameter traces, i.e. the 1D functions of parameter value vs. step number for each parameter separately:

```python
fig,axes = plt.subplots(len(p0), 1, figsize=(5,7), sharex=True)

for i in range(len(p0)):
    axes[i].plot(chain[:,i], marker='', drawstyle='steps')

axes[0].axhline(a_true, color='r', label='true')
axes[0].legend(loc='best')
axes[0].set_ylabel('$a$')

axes[1].axhline(b_true, color='r')
axes[1].set_ylabel('$b$')

fig.tight_layout()
```

From these trace plots, we can see by eye that it takes the sampler about a few hundred steps to converge. When we look at the samples returned or when we compute our summary statistics, we don't want to include these parameter values! In addition, there is likely some correlation between nearby steps. We can attempt to remove some of the correlated steps by _thinning_ the chain, i.e. by downsampling. We can do both simultaneously using Python indexing tricks. Certainly by step 2000 the chains look converged, so from there on we'll keep only every 8th step:

```python
good_samples = chain[2000::8]
good_samples.shape
```

We're left with 774 samples; we hope these are approximately uncorrelated, converged samples from the posterior pdf (there are other ways we can check, but these are out of scope for this workshop). Now you have to choose what summary statistics to report. You have some options, but a reasonable choice is to report the median, 16th, and 84th percentiles:

```python
low,med,hi = np.percentile(good_samples, [16, 50, 84], axis=0)
upper, lower = hi-med, med-low

disp_str = ""
for i,name in enumerate(['a', 'b']):
    fmt_str = '{name}={val:.2f}^{{+{plus:.2f}}}_{{-{minus:.2f}}}'
    disp_str += fmt_str.format(name=name, val=med[i], plus=upper[i], minus=lower[i])
    disp_str += r'\quad '

from IPython import display
disp_str = "${}$".format(disp_str)
display.Latex(data=disp_str)
```

Recall that the true values are:

```python
a_true, b_true
```

We've now done this problem the Bayesian way as well! Now, instead of drawing the "best-fit" line over the data, we can take a handful of samples and plot a line for each of the samples, as a way to visualize the uncertainty we have in the model parameters:

```python
plt.figure(figsize=(6,5))
plt.errorbar(x, y, y_err, linestyle='none', marker='o', ecolor='#666666')

x_grid = np.linspace(x.min()-0.1, x.max()+0.1, 128)
for pars in good_samples[:128]: # only plot 128 samples
    plt.plot(x_grid, line_model(pars, x_grid),
             marker='', linestyle='-', color='#3182bd', alpha=0.1, zorder=-10)

plt.xlabel('$x$')
plt.ylabel('$y$')
plt.tight_layout()
```

Or, we can plot the samples using a _corner plot_ to visualize the structure of the 2D and 1D (marginal) posteriors:

```python
# uncomment and run this line if the import fails:
# !source activate statsseminar; pip install corner
import corner
```

```python
fig = corner.corner(chain[2000:], bins=32, labels=['$a$', '$b$'], truths=[a_true, b_true])
```

---

## Finally, the problem you came here for: fitting a straight line to data with intrinsic scatter

We made it! We're now ready to do the problem we set out to do. In the initial model, we assumed that we knew the uncertainties in our measurements exactly and that the data were drawn from a one-dimensional line. We're now going to relax that assumption and assume that either (a) the data uncertainties have been underestimated or (b) there is intrinsic scatter in the true model (in the absence of other information, these two ideas are degenerate). Let's first generate some data. We'll assume the latter of the two ideas, and we'll further assume that the model line is convolved with an additional Gaussian in the $y$ direction, with the new parameter being the intrinsic width of the relation expressed as a variance $V$:

```python
V_true = 0.5**2
n_data = 42

# we'll keep the same parameters for the line as we used above
```

```python
x = rnd.uniform(0, 2., n_data)
x.sort() # sort the values in place

y = a_true*x + b_true

# Heteroscedastic Gaussian uncertainties only in y direction
y_err = rnd.uniform(0.1, 0.2, size=n_data) # randomly generate uncertainty for each datum

# add Gaussian intrinsic width
y = rnd.normal(y, np.sqrt(y_err**2 + V_true)) # re-sample y data with noise and intrinsic scatter
```

```python
plt.errorbar(x, y, y_err, linestyle='none', marker='o', ecolor='#666666')
plt.xlabel('$x$')
plt.ylabel('$y$')
plt.tight_layout()
```

Let's first naively fit the data assuming no intrinsic scatter using least-squares:

```python
X = np.vander(x, N=2, increasing=True)
Cov = np.diag(y_err**2)
Cinv = np.linalg.inv(Cov)
```

```python
best_pars = np.linalg.inv(X.T @ Cinv @ X) @ (X.T @ Cinv @ y)
pars_Cov = np.linalg.inv(X.T @ Cinv @ X)
```

```python
plt.errorbar(x, y, y_err, linestyle='none', marker='o', ecolor='#666666')

x_grid = np.linspace(x.min()-0.1, x.max()+0.1, 128)
plt.plot(x_grid, line_model(best_pars[::-1], x_grid), marker='', linestyle='-', label='best-fit line')
plt.plot(x_grid, line_model([a_true, b_true], x_grid), marker='', linestyle='-', label='true line')

plt.xlabel('$x$')
plt.ylabel('$y$')

plt.legend(loc='best')
plt.tight_layout()
```

The covariance matrix for the parameters is:

```python
pars_Cov
```

We clearly get a biased result and yet _very_ precise measurements of the parameters when we don't take in to account the intrinsic scatter. What we need to do now is modify out model to include the scatter as a free parameter. Unfortunately, it enters the model non-linearly so there is no solution using linear algebra or least-squares. Instead, we just write a new likelihood function and optimize it numerically. One choice we'll make is to use the parameter $\ln{V}$ instead of $V$ for reasons I'll explain later. To implement the new model, we'll subclass our `StraightLineModel` class and define new likelihood and prior functions.

### Exercise 7:

Subclass the `StraightLineModel` class and implement new prior and likelihood functions (`ln_prior` and `ln_likelihood`). The our model will now have 3 parameters: `a`, `b`, and `lnV` the log of the intrinsic scatter variance. Use flat priors on all of these parameters. In fact, we'll be even lazier and forget the constant normalization terms: if a parameter vector is within the ranges below, return 0. (log(1.)) otherwise return -infinity:

#### Solution:

```python
class StraightLineIntrinsicScatterModel(StraightLineModel):

    def ln_prior(self, pars):
        """ The prior only depends on the parameters """

        a, b, lnV = pars

        # flat priors on a, b, lnV
        if a < -10 or a > 10 or b < -100. or b > 100. or lnV < -10. or lnV > 10.:
            return -np.inf

        # this is only valid up to a numerical constant
        return 0.

    def ln_likelihood(self, pars):
        """ The likelihood function evaluation requires a particular set of model parameters and the data """
        a,b,lnV = pars
        V = np.exp(lnV)

        N = len(self.y)
        dy = self.y - line_model([a,b], self.x)
        ivar = 1 / (self.y_err**2 + V) # inverse-variance now includes intrinsic scatter
        return -0.5 * (N*np.log(2*np.pi) - np.sum(np.log(ivar)) + np.sum(dy**2 * ivar))
```

```python
scatter_model = StraightLineIntrinsicScatterModel(x, y, y_err)
```

```python
x0 = [5., 5., 0.] # starting guess for the optimizer

# we have to minimize the negative log-likelihood to maximize the likelihood
result_ml_scatter = minimize(lambda *args: -scatter_model.ln_likelihood(*args),
                             x0=x0, method='BFGS')
result_ml_scatter
```

```python
plt.errorbar(x, y, y_err, linestyle='none', marker='o', ecolor='#666666')

x_grid = np.linspace(x.min()-0.1, x.max()+0.1, 128)
plt.plot(x_grid, line_model(result_ml_scatter.x[:2], x_grid), marker='', linestyle='-', label='best-fit line')
plt.plot(x_grid, line_model([a_true, b_true], x_grid), marker='', linestyle='-', label='true line')

plt.xlabel('$x$')
plt.ylabel('$y$')

plt.legend(loc='best')
plt.tight_layout()
```

```python
V_true, np.exp(result_ml_scatter.x[2])
```

It looks like the maximum likelihood estimate is a little bit better, and we get a reasonable measurement of the intrinsic scatter, but none of this gives us a handle on the uncertainty. How do we quantify the uncertainty in the now 3 parameters? We'll just run MCMC.

### Exercise 8:

To quantify our uncertainty in the parameters, we'll run MCMC using the new model. Run MCMC for 65536 steps and visualize the resulting chain. Make sure the acceptance fraction is between ~25-50%.

#### Solution:

```python
p0 = [6., 6., -1.]
chain,_,acc_frac = run_metropolis_hastings(p0, n_steps=2**16, model=scatter_model,
                                           proposal_sigmas=[0.15,0.15,0.2])
acc_frac
```

```python
fig,axes = plt.subplots(len(p0), 1, figsize=(5,7), sharex=True)

for i in range(len(p0)):
    axes[i].plot(chain[:,i], marker='', drawstyle='steps')

axes[0].axhline(a_true, color='r', label='true')
axes[0].legend(loc='best')
axes[0].set_ylabel('$a$')

axes[1].axhline(b_true, color='r')
axes[1].set_ylabel('$b$')

axes[2].axhline(np.log(V_true), color='r')
axes[2].set_ylabel(r'$\ln V$')

fig.tight_layout()
```

```python
fig = corner.corner(chain[2000:], bins=32, labels=['$a$', '$b$', r'$\ln V$'],
                    truths=[a_true, b_true, np.log(V_true)])
```

Now we'll again compute the percentiles for the 1D, marginal distributions:

```python
good_samples = chain[2000::8]
good_samples.shape
```

```python
low,med,hi = np.percentile(good_samples, [16, 50, 84], axis=0)
upper, lower = hi-med, med-low

disp_str = ""
for i,name in enumerate(['a', 'b', r'\ln V']):
    fmt_str = '{name}={val:.2f}^{{+{plus:.2f}}}_{{-{minus:.2f}}}'
    disp_str += fmt_str.format(name=name, val=med[i], plus=upper[i], minus=lower[i])
    disp_str += r'\quad '

disp_str = "${}$".format(disp_str)
display.Latex(data=disp_str)
```

Compare this to the diagonal elements of the covariance matrix we got from ignoring the intrinsic scatter and doing least-squares fitting:

```python
disp_str = ""
for i,name in zip([1,0], ['a', 'b']):
    fmt_str = r'{name}={val:.2f} \pm {err:.2f}'
    disp_str += fmt_str.format(name=name, val=best_pars[i], err=np.sqrt(pars_Cov[i,i]))
    disp_str += r'\quad '

disp_str = "${}$".format(disp_str)
display.Latex(data=disp_str)
```

The parameter uncertainties estimated from the MCMC samples are much larger -- this reflects our uncertainty about the intrinsic scatter of the points. Precision is highly model dependent.


# Nearly-incompressible flow: the Tammann equation of state

In this chapter, we consider the one-dimensional Riemann problem for the Euler equations with the Tammann equation of state (EOS). The Tammann EOS, also known as stiffened gas EOS, is used to model almost-incompressible fluids. It can be physically understood as modeling an ideal gas under very high ambient pressures. The Riemann problem we want to solve consists of the one-dimensional Euler equations,

\begin{align*}
    \left[\begin{array}{c} \rho \\   \rho u \\    E   \end{array} \right]_t +
    \left[\begin{array}{c} \rho u \\   \rho u^2 + p \\     u(E+p)   \end{array} \right]_x
=   0,
\end{align*}
where $\rho$ is density, $u$ the velocity, $E$ the internal energy, and $p$ the pressure change relative to an ambient pressure $p_{\infty}$. As usual, the subcripts $x,t$ denote partial derivatives. The initial conditions are given by the left and right constant states $q_l=[\rho_l, \rho_l u_l, E_l]$ and $q_r=[\rho_r, \rho_r u_r, E_r]$. The Tamann EOS is given by

\begin{align*}
p (\rho,e)=\rho e(\gamma -1) - \gamma p_{\infty},
\end{align*}
where $e$ is the specific internal energy. For a general equation of state $p (\rho,e)$, the speed of sound is given by the isentropic derivative of the pressure with respect to density (keeping entropy constant),

\begin{align*}
  c = \sqrt{\frac{\partial p(\rho,e)}{\partial \rho} \bigg|_s} = \sqrt{\frac{\partial p(\rho,e)}{\partial \rho} + \frac{p(\rho,e)}{\rho^2}\frac{\partial p(\rho,e)}{\partial e}}.
\end{align*}
In the case of the Tammann EOS, we obtain

\begin{align*}
 c= \sqrt{\gamma\frac{p+p_{\infty }}{\rho}}.
\end{align*}
Note that this is almost the same as the [sound speed in an ideal gas](Euler_equations.ipynb#Characteristic-velocities) except for the additional $p_{\infty}$ term. This means that the fluid behaves similarly to an ideal gas that is already at pressure $p_{\infty }$. When $p_{\infty } \gg 1$, the Tammann EOS emulates an ideal gas at a very high pressure, which can be used to model almost-incompressible fluids like water. They are almost-incompressible because a very large change is pressure is required to produce even a small change in density.  Note the sound speed (as well as the density and any other physical quantity) is well defined for negative values of $p$, with the only limitation that $p+p_{\infty}>0$. We will revisit this issue in [the last section of this chapter](#Strong-rarefaction-waves).

Using the equation of state, the state of the system can also be written in the primitive variables $[\rho, u, p]$.

## Exact solution

The exact solution to the Riemann problem with the Tammann EOS when the fluid parameters $\gamma$ and $p_\infty$ do not vary can be found in (Ivings & Toro 1998). For the more general case when the left and right states may have different values of $\gamma$ and $p_\infty$, the solution can be found in the supplemental material of the work (del Razo, 2017). In this chapter we will focus on the general case. The presentation follows that of (del Razo, 2017).

Just as in [the earlier chapter on ideal gas flow](Euler_equations.ipynb), the solution of the Riemann problem considered here consists of three waves.  The 1-wave and 3-wave are genuinely nonlinear (rarefactions or shocks), while the 2-wave is a linearly degenerate contact discontinuity. The solution thus includes four constant regions with solutions $q_l, q_{*l}, q_{*r}, q_r$, separated by the three waves. Just as in the case of the ideal gas, the velocity and pressure do not change across the contact wave, so we denote them simply by $u_*$ and $p_*$. Note that the interface between the fluids that were initially to the left and right of $x=0$ moves along the contact discontinuity, so if the fluid parameters $\gamma$ and $p_\infty$ vary between the left and right states then the this jump in parameters also moves with the contact discontinuity.  Consequently, the 1-wave propagates through left-state fluid and the 3-waves propagates through right-state fluid.  (This will be true even in the supersonic case, where all the wave speeds have the same sign.) This allows us to determine the Hugoniot loci and integral curves for these two nonlinear waves using the fact that fluid parameters $\gamma$ and $p_\infty$ are always constant across each of these waves.

In order to determine whether the 1-wave and 3-wave are rarefactions or shocks, we will need to create a function of the
middle state pressure $p_*$ that ensures that the velocity $u_*$ across the contact discontinuity is consistent. Since we know that the left state velocity $u_l$ should be connected by a rarefaction or shock to $u_*$, we write $u_*=u_l + [u]_1$, where $[u]_1$ is the jump of the velocity across the 1-wave. In a similar manner, we also
know the 3-wave should be a shock or rarefaction, so we write $u_*= u_r - [u]_3$.  Therefore, it is useful to define

\begin{align}
  \phi_l(p) &= u_l - \mathcal{F}_l(p),   \label{CDvel-1}  \\
  \phi_r(p) &= u_r + \mathcal{F}_r(p).   \label{CDvel-2}
\end{align}

where the value of $\mathcal{F}_{l,r}(p) = -[u]_{1,3}$ depends on whether the wave in question is a shock or a rarefaction (signs were
chosen for notational consistency). As we expect these two equations to yield the same contact discontinuity
velocity $u_* = \phi_l(p_*) = \phi_r(p_*)$, we have

\begin{align}
 \Phi(p_*)= \phi_r(p_*) - \phi_l(p_*) = 0.
 \label{Phi}
\end{align}

This nonlinear equation will yield the pressure $p_*$ that provides consistency between the type of waves (rarefactions or shocks), their speeds and the contact discontinuity velocity $u_*$. As we mentioned before, the shape of $\phi_k(p_*)$ will depend on whether the states are connected by a shock wave or rarefaction. Once
the $p_*$ has been found, the contact discontinuity velocity can be found from the expressions just derived. However, it is not yet clear how to calculate the density or the speeds of the 1-wave
and 3-wave. It will become obvious how to obtain these quantities once we write the explicit equations for our system further below.

Before coding the exact solution, we should first note that having a rarefaction or shock in the 1-wave
and 3-wave will depend on the pressure $p_*$, in [the same way it did for an ideal gas](Euler_equations.ipynb#Entropy-condition).  If the pressure is higher on the side toward which the wave is propagating (i.e., in front of the wave), it must be a rarefaction. If the pressure is lower in front of the wave, it must be a shock. In the Euler equations, this yields four possible cases for the value $\Phi(p_*)$, just as in the solution with an ideal gas EOS:

 * 1-rarefaction, 3-rarefaction: $p_*< p_l$ and $p_*<p_r$
 $\Phi(p_*)= \phi_r^R(p_*) - \phi_l^R(p_*)$,
 * 1-shock,       3-rarefaction $p_l \le p_* \le p_r$
 $\Phi(p_*)= \phi_r^R(p_*) - \phi_l^S(p_*)$,
 * 1-rarefaction,       3-shock $p_r \le p_* \le p_l$
 $\Phi(p_*)= \phi_r^S(p_*) - \phi_l^R(p_*)$,
 * 1-shock, 3-shock: $p_* > p_l$ and $p_*>p_r$
 $\Phi(p_*)= \phi_r^S(p_*) - \phi_l^S(p_*)$,

where the index $S,R$ indicates whether the function $\phi$ was obtained by using the Rankine-Hugoniot equations to connect states by shocks or the Riemann invariants to connect them by rarefactions respectively.

We now derive the functions $\phi_{k}^{\mu}$ for all four cases with $k=l,r$ (left and right) and $\mu=R,S$ (rarefaction and shock).  We also show how to obtain the density and the missing wave speeds. We write $s_n$ for the speed of the wave belonging to family $n$ ($n=1,2,3$).

### Rankine-Hugoniot conditions for shock waves
Let us first consider the case in which both the 1-wave and the 3-wave are shocks. The velocity of each
shock will be given by the Rankine-Hugoniot conditions. These conditions will also allow us to calculate $\phi_k^S(p_*)$ with $k=l,r$.  The following analysis is written to be applicable to both the 1-wave
velocity $s_1$ and the 3-wave velocity $s_3$, by considering $s_n$ and $q_k$, with $n=1, k=l$ for the 1-wave and $s=3, k=r$ for the 3-wave.

The Rankine-Hugoniot conditions for
the Euler equations can be written as (see e.g. (Ivings & Toro 1998)),

\begin{align}
 \rho_k \omega_k = \rho_{*k} \omega_*, \label{RH-mass} \\
 \rho_k \omega_k^2 + p_k = \rho_{*k} \omega_*^2 + p_{*k}, \label{RH-mom}  \\
 \frac{1}{2} \omega_k^2 + h_k = \frac{1}{2} \omega_*^2 + h_{*k}, \label{RH-ene}
\end{align}
where $\omega_k = u_k - s_n$ , $\omega_* = u_* - s_n$, and the specific enthalpy is given
by $h = e+(p+p_\infty)/\rho$ with $e$ the specific internal energy that relates to the internal energy
of our original variables by $E = \rho e + \rho u^2/2$. We will use these relations to find $\phi_k^S(p_*)$ and the wave speeds $s_n$.

#### Finding $\phi_l^S(p_*)$ and $\phi_r^S(p_*)$ and $s_1$ and $s_3$:
We start by defining the mass fluxes $F_k$ for $k=l,r$ as

\begin{align}
 F_l = \rho_l \omega_l = \rho_{*l} \omega_* \label{Ql} \\
 F_r = -\rho_r \omega_r = -\rho_{*r} \omega_*. \label{Qr}
\end{align}

As $\omega_k = u_k - s_n$, from these two equations we can obtain the wave speeds in terms of $F_l$ and $F_r$,

\begin{align}
 s_1 = u_l - \frac{F_l}{\rho_l}, \  \  \
 s_3 = u_r + \frac{F_r}{\rho_r}. \label{RH-speeds}
\end{align}

We still need to find $F_l$ and $F_r$, so we substitute equations (\ref{Ql}) and (\ref{Qr})
into (\ref{RH-mom}) to obtain

\begin{align}
  F_l &= \frac{\tilde{p}_{*l} -\tilde{p}_l}{\omega_l - \omega_*} = \frac{p_* -p_l}{u_l - u_*} \label{Ql-1} \\
  F_r &= -\frac{\tilde{p}_{*r} -\tilde{p}_r}{\omega_r - \omega_*} = -\frac{p_* -p_r}{u_r - u_*}, \label{Qr-1}
\end{align}
where $\tilde{p}_\kappa = p_\kappa + p_{\infty \kappa}$ is defined to simplify future notation
with $\kappa = l, *l, *r$ and $r$. Note that $\tilde{p}_{*l} \neq \tilde{p}_{*r}$ and
that $\tilde{p}_{*k} -\tilde{p}_k=p_* -p_k$ since $p_*=p_{*l}=p_{*r}$ and $p_{\infty k} = p_{\infty *k}$  ($k=l,r)$.
Solving for $u_*$ we obtain the equations

\begin{align}
 u_* = u_l - \frac{\tilde{p}_{*l} - \tilde{p}_l}{F_l} = \phi_l^S(p_*) \\
 u_* = u_r + \frac{\tilde{p}_{*r} - \tilde{p}_r}{F_r} = \phi_r^S(p_*)
 \label{RH-phis}
\end{align}

Comparing to equations (\ref{CDvel-1}, \ref{CDvel-2}), we notice $\mathcal{F}_k(p_*) = \frac{p_* - p_k}{F_k}$. We also notice we have almost
obtained the $\phi$ functions we are looking for, though we still need to find $F_l$ and $F_r$ in terms of known
variables.

#### Finding $F_l$ and $F_r$:

From equations (\ref{Ql}) and  (\ref{Ql-1}), we know
that

\begin{align*}
\frac{\tilde{p}_{*l} -\tilde{p}_l}{\omega_l - \omega_*} = \rho_l \omega_l.
\end{align*}

Solving for $\omega_*$, substituting the solution into (\ref{Ql}) and subtituting the $w_l$ for $F_l/\rho_l$,
we obtain a new equation that we can solve for $F_l$ that yields

\begin{align}
 F_k = \sqrt{\rho_k \rho_{*k} \frac{\tilde{p}_{*k} -\tilde{p}_k}{\rho_{*k} - \rho_k}}. \label{Qk}
\end{align}
with $k=l,r$, since we repeated the same process for $F_r$ and obtained exactly the same equation.
However, we still don't know $\rho_{*k}$, for this we will need our third Rankine-Hugoniot condition (\ref{RH-ene}).

#### Finding $\rho_{*l}$ and $\rho_{*r}$:

From equation (\ref{RH-ene}), we can obtain

\begin{align}
 h_{*k} - h_k &= \frac{1}{2}\left(w_k^2 - w_*^2\right) \nonumber\\
              &= \frac{1}{2}\left(\pm \frac{F_k^2}{\rho_k^2} \mp \frac{F_k^2}{\rho_{*k}^2}\right) \nonumber \\
              &= \frac{1}{2}\left(\frac{1}{\rho_k} + \frac{1}{\rho_{*k}}\right)(\tilde{p}_{*k}-\tilde{p}_k), \label{delta-h}
\end{align}

where the sign above is used for $k=l$ and the one below for $k=r$, and we used equations
(\ref{Ql}) and (\ref{Qr}) for the second line and (\ref{Qk}) for the third line. We can now substitute
the specific enthalpy $h=\gamma \tilde{p}/(\rho(\gamma -1))$
in equation (\ref{delta-h}) to obtain

\begin{align*}
 \frac{\gamma_{*k}}{\gamma_{*k} - 1}\frac{\tilde{p}_{*k}}{\rho_{*k}} - \frac{\gamma_{k}}{\gamma_{k} - 1}\frac{\tilde{p}_k}{\rho_{k}}
 =\frac{1}{2}\left(\frac{1}{\rho_k} + \frac{1}{\rho_{*k}}\right)(\tilde{p}_{*k}-\tilde{p}_k).
\end{align*}

As the interface is the contact discontinuity, the jump in the parameters is only across the contact discontinuity, so
$\gamma_{*k}=\gamma_k$. Now we can solve for the unknown density:

\begin{align}
 \rho_{*k} = \rho_k \left(\frac{\frac{\tilde{p_*}}{\tilde{p_k}} +
 \frac{\gamma_k-1}{\gamma_k+1}}{\frac{\tilde{p_*}}{\tilde{p_k}}\frac{\gamma_k-1}{\gamma_k+1}+1}\right).
 \label{rho-k}
\end{align}

Substituting this result into equation (\ref{Qk}), we obtain $F_k$ in terms of $p_*$ and known variables:

\begin{align}
  F_k = \sqrt{\rho_k \frac{\tilde{p}_{*k} + \tilde{p}_{k}\frac{\gamma_k - 1}{\gamma_k +1}}{\frac{2}{\gamma_k+1}}}.
\label{Qk-f}
\end{align}

Subsitituting (\ref{Qk-f}) into (\ref{RH-phis}), we can calculate the nonlinear functions $\phi^S_{l,r}(p_*)$
in terms of known variables, which yields

\begin{align}
\phi_l^S(p_*) = u_l - (p - p_l)\sqrt{\frac{\alpha_l}{p_* + p_{\infty l} + \beta_l}} \\
\phi_r^S(p_*) = u_r + (p - p_r)\sqrt{\frac{\alpha_r}{p_* + p_{\infty r} + \beta_r}}
\end{align}
with $\alpha_k = 2.0/((\gamma_k + 1.0)\rho_l)$ and $\beta_k = (p_l + p_{\infty l})(\gamma_l - 1.0) / (\gamma_l + 1.0)$. The functions $\phi_k^S$ define the Hugoniot Loci, and they allow us to construct equation (\ref{Phi}) and solve it using a Newton method or other root finder in order to obtain the value of $p_*$.

### Riemann invariants for rarefaction waves

Riemann invariants are variables that remain constant through simple waves such as centered rarefactions. The Riemann invariants across
the 2-wave (contact discontinuity) are the pressure $p_*$ and the normal velocity $u_*$. The Riemann invariants for the 1-wave and 3-wave are
the entropy and the quantities

\begin{align}
 u_l+\frac{2c_l}{\gamma_l-1} = u_* + \frac{2c_{*l}}{\gamma_l-1}, \label{RI1} \\
 u_r-\frac{2c_r}{\gamma_r-1} = u_* - \frac{2c_{*r}}{\gamma_r-1}. \label{RI2}
\end{align}
The speed of sound $c_k$ with the Tammann EOS is given by

\begin{align}
 c_k= \sqrt{\gamma_k\frac{p_k+p_{\infty k}}{\rho_k}}.  \label{sndsp}
\end{align}

As the entropy is invariant along the rarefaction wave, we can use the Tammann EOS isentropic relation to obtain the density in the middle states,

\begin{align}
 \rho_{*k} = \rho_k \left(\frac{\tilde{p}_{*k}}{\tilde{p}_k}\right)^{1/\gamma}.  \label{isentrop}
\end{align}

Solving (\ref{RI1}) and (\ref{RI2}) for $u_*$ and using equations (\ref{sndsp}) and (\ref{isentrop}), we immediately obtain

\begin{align*}
 u_* = u_l + \frac{2c_l}{\gamma_l-1} \left[
 1-\left(\frac{\tilde{p}_{*l}}{\tilde{p}_l}\right)^{\frac{\gamma_l-1}{2\gamma_l}}
 \right] = \phi_l^R(p_*), \\
 u_* = u_r - \frac{2c_r}{\gamma_r-1} \left[
 1-\left(\frac{\tilde{p}_{*r}}{\tilde{p}_r}\right)^{\frac{\gamma_r-1}{2\gamma_r}}
 \right] = \phi_r^R(p_*).
\end{align*}

These functions define the integral curves, and they indicate which pairs of states can be connected by a rarefaction. Note that since $\tilde{p}_\kappa = p_\kappa + p_{\infty \kappa}$, when we compare to equations (\ref{CDvel-1}) and (\ref{CDvel-2}), we obtain
the functions $\phi_{l,r}^R$. The rarefaction head velocities are given by $u_l-c_l$
and $u_r+c_r$; the tail velocities are $u_*-c_{*l}$ and $u_*+c_{*r}$.

In order to compute the complete structure of the rarefaction wave (Ivings & Toro 1998,  LeVeque 2002), we can use the
Riemann invariants from Eqs. \eqref{RI1} and \eqref{RI2}, along with Eq. \eqref{sndsp} and the isentropic
relation from Eq. \eqref{isentrop}. The solution for the 1-rarefaction wave along the rays $x/t=\xi= u_{rar1} - c_{rar1}$ is then

\begin{align*}
 u_{rar1}(\xi) &= \frac{u_l(\gamma_l - 1) + 2(\xi + c_l)}{\gamma_l + 1}, \\
 \rho_{rar1}(\xi) &= \rho_l \left[\frac{u_{rar1}(\xi) - \xi}{c_l}\right]^{\frac{2}{\gamma_l -1}}, \\
 p_{rar1}(\xi) &= \tilde{p}_l \left[\frac{u_{rar1}(\xi) - \xi}{c_l}\right]^{\frac{2\gamma_l}{\gamma_l -1}} - p_{\infty l},
\end{align*}

and for a 3-rarefaction wave along the rays $x/t=\xi= u_{rar3} + c_{rar3}$ is,

\begin{align*}
 u_{rar3}(\xi) &= \frac{u_r(\gamma_r - 1) + 2(\xi - c_r)}{\gamma_r + 1}, \\
 \rho_{rar3}(\xi) &= \rho_r \left[\frac{\xi - u_{rar3}(\xi)}{c_r}\right]^{\frac{2}{\gamma_r -1}}, \\
 p_{rar3}(\xi) &= \tilde{p}_r \left[\frac{\xi - u_{rar3}(\xi)}{c_r}\right]^{\frac{2\gamma_r}{\gamma_r -1}} - p_{\infty r}.
\end{align*}

Now that we know the integral curves given by $\phi^{R}_{l,r}$ (for the rarefactions), we can construct the function $\Phi(p_*)$ from (\ref{Phi}) for any of the 4 possible scenarios.

\begin{align*}
\Phi(p_*) & = \begin{cases}
 \phi_r^R(p_*) - \phi_l^R(p_*) & \text{if } p_* < p_l \text{ and } p_*<p_r \\
 \phi_r^R(p_*) - \phi_l^S(p_*) & \text{if } p_l \le p_* \le p_r \\
 \phi_r^S(p_*) - \phi_l^R(p_*) & \text{if } p_r \le p_* \le p_l \\
 \phi_r^S(p_*) - \phi_l^S(p_*) & \text{if } p_* > p_l \text{ and } p_* > p_r
              \end{cases}
\end{align*}

The value of $p_*$ will be
found by numerically finding the root of $\Phi(p_*)=0$. Once $p_*$ is found, $u_*$, $\rho_{*l}$, $\rho_{*r}$, $s_1$
and $s_3$ can be found using the relations above. As we
know the three wave speeds $s_1$, $s_2$ and $s_3$ and the values of the primitive variables $[\rho,u,p]$ in all the 4 states, we have solved the Riemann problem.

### Solution in the pressure-velocity plane

As we did in [the chapter on ideal gas flow](Euler_equations.ipynb), here we visualize the solution in the plane whose axes are pressure and velocity. We have a left and a right state such that each one is connected to some middle state by a shock or a rarefaction, depending on the entropy conditions. The connection is established by curves in the phase plane: Hugoniot loci in the case of shocks and integral curves in the case of rarefactions. The intersection between these curves yields the middle state, which is the most important element of the solution to the Riemann problem. Let's look at an interactive plot illustrating the solution in the phase plane. The initial states and the middle state are indicated as black points; the Hugoniot loci are plotted in red and the integral curves in blue. The physical solution satisfying the entropy conditions is plotted as a continous line, whereas the unphysical solution is shown as a dashed line.

```python
%matplotlib inline
import numpy as np
from ipywidgets import interact
from ipywidgets import widgets
import matplotlib.pyplot as plt
from exact_solvers import euler_tammann, interactive_pplanes
from utils import riemann_tools
import warnings
warnings.filterwarnings("ignore")
```

```python
# Initial states [density, velocity, pressure]
ql = [600.0, 10.0, 50000.0]
qr = [50.0, -10.0, 25000.0]
# Tammann EOS parameters [gamma, pinf]
paramsl = [1.4, 0.0]
paramsr = [7.0, 100.0]

interactive_pplanes.euler_tammann_interactive_phase_plane(ql,qr,paramsl,paramsr)
```

### Full solution
In order to get the full solution, we need some additional quantities. Equations (\ref{RH-phis}) will yield the contact discontinuity speed $s_2=u_*$ in terms of $p_*$. We can also calculate $F_l$ and $F_r$ from (\ref{Qk-f}), and we can substitute in (\ref{RH-speeds}) to obtain the corresponding wave speeds. With this, we can provide the full solution of the Riemann problem. In the next section, we will show the full solution for different examples.   we define a function `plot_riemann_solution` to plot the exact solution of the Riemann problem and parameters that we will use for the different examples.

```python
def plot_riemann_solution(q_l, q_r, auxl, auxr):
    ex_states, ex_speeds, reval, wave_types, varsout = \
        euler_tammann.exact_riemann_solution(q_l ,q_r, auxl, auxr,
                                             varin='primitive',
                                             varout='primitive')
    varnames = euler_tammann.primitive_variables
    plot_function = \
            riemann_tools.make_plot_function(ex_states, ex_speeds,
                                             reval, wave_types,
                                             layout='vertical',
                                             variable_names=varnames,
                                             aux=(np.array(auxl), np.array(auxr)),
                                             plot_chars=[euler_tammann.lambda1,
                                                         euler_tammann.lambda2,
                                                         euler_tammann.lambda3],
                                            contact_index=1)

    dropdown = widgets.Dropdown(options=[None,1,2,3],
                                description='Show characteristics')
    return interact(plot_function,
                    t=widgets.FloatSlider(value=0.0,min=0,max=1.0),
                    which_char=dropdown);
```

## Examples of Riemann Solutions

### Problem 1: Sod shock tube in water
We first consider the classic Sod shock tube problem with higher density and pressure on the left state than on the right state and both sides initially at rest. The main difference with the classic problem is that in this case the shock tube is filled with water and not with air. The propagation of waves in water can be accurately modeled using $\gamma = 7.15$, $p_{\infty} = 3\times 10^8$ and $\rho\approx 1000 kg/m^3$; see (del Razo 2016). A value of $p_{atm}=101325 Pa$ is used for Atmospheric pressure.

```python
gamma_water = 7.15
gamma_air = 1.4
pinf_water = 3.e8
pinf_air = 0.
patm = 101325.0
```

We take a very large jump in pressure so that it's clear the left-going wave is a rarefaction. With a more modest jump the rarefaction wave barely spreads and appears more like a discontinuity.

```python
q_l = [1500.0, 0.0, 3000*patm]
q_r = [1000.0, 0.0, patm]
gamma = [gamma_water, gamma_water]
pinf = [pinf_water, pinf_water]

plot_riemann_solution(q_l, q_r, [gamma[0],pinf[0]], [gamma[1],pinf[1]]);
```

Note the small density difference accross left-going wave and the very compressed rarefaction fan
are consequence of the low compressibility of water.

### Problem 2: Transmission and reflection at an air-water interface
We can use the Euler equations with the Tammann EOS to model wave propagation in almost-incompressible media, such as water. Since we have the solution for different EOS parameters in the left and right states, we can model interaction of different fluids across the interface. For instance, we can use this model to study the Riemann problem across an air-water interface. Air is modeled by the ideal gas EOS, which corresponds to $\gamma=1.4$ and $p_{\infty} = 0$, and water can be modeled using the same parameters as the example above.

Consider the problem of a shock traveling from left to right through the air and hitting the interface at time 0. Then this wave should be partly transmitted as a shock wave moving to the right in the water, and partly reflected as a lef-going wave in the air.  If we suppose the shock wave hits the interface at time $t=0$, then we can calculate the transmitted and reflected waves by solving the Riemann problem starting at $t=0$ for a case with a positive pressure and velocity in the left state and ambient water conditions in the right state.

As we want to model an air-water interface, the left state has only air while the right state has only water with their corresponding densities. Although in a realistic two or three dimensional setting, the water would spill into the left state, we can suppose that the water is being held in place by a container with a thin wall. In this setting, the current model has been shown to be a valid approximation (del Razo 2017). The solution has the following form:

```python
q_l = [1.0, 350.0, 300*patm] # Initial condition for shock in air
q_r = [1000.0, 0.0, patm] # Initial condition in still water
gamma = [gamma_air, gamma_water]
pinf = [pinf_air, pinf_water]

plot_riemann_solution(q_l, q_r, [gamma[0],pinf[0]], [gamma[1],pinf[1]]);
```

Note we exaggerated again the initial jump in pressure, so the jump in the density and velocity of the right-going wave could be appreciated in the plots; otherwise, it seems there is no jump in the right-going wave.

### Problem 3: Water-air interface
Here we consider a setup similar to that of Problem 2, but we invert the materials. The left region has water, while the right region has air.

```python
q_l = [1000.0, 350.0, 300*patm] # Initial condition for shock in air
q_r = [1.0, 0.0, patm] # Initial condition in still water
gamma = [gamma_water, gamma_air]
pinf = [pinf_water, pinf_air]

plot_riemann_solution(q_l, q_r, [gamma[0],pinf[0]], [gamma[1],pinf[1]]);
```

Once again we exaggerated the jump in pressure to be able to observe the jumps of the solution. Nonetheless, the density jump in the right-going wave is still not perceptible. A more exaggerated initial condition can make this jump more evident.

## Strong rarefaction waves

As we mentioned when we introduced the Tammann EOS, the relevant physical quantities remain well-defined for
negative values of $p$, as long as $p+p_{\infty}>0$. For instance, the densities \eqref{rho-k}, \eqref{isentrop} and the sound speed \eqref{sndsp} are all well defined in this case. This is consistent with the interpretation of the Tammann EOS as a shift of the zero pressure point of an ideal gas.

As an example of this scenario, we consider a strong symmetric expansion wave in water. Analogous to the Euler equations using an ideal gas EOS, we consider equal densities and pressures, and equal and opposite velocities, with the initial states moving away from each other. The result is two rarefaction waves (the contact has zero strength). Notice that the value in the pressure plot is negative; it reaches a value of about $-2.8\times 10^8$, so $p+p_{\infty}$ is still positive.  Of course, it is questionable whether the model remains accurate when $p+p_{\infty}$ is nearly zero, as it is in this example.

```python
velocity = 350
q_l = [1000.0, -velocity, 2*patm]
q_r = [1000.0, velocity, 2*patm]
auxl = [gamma_water, pinf_water]
auxr = [gamma_water, pinf_water]
ex_states, ex_speeds, reval, wave_types, varsout = \
        euler_tammann.exact_riemann_solution(q_l ,q_r, auxl, auxr,
                                             varin='primitive',
                                             varout='primitive')

print('Middle-state pressure p = %10.3e' % ex_states[2,1])
print('        p_inf for water = %10.3e,' % pinf_water)
print('                p+p_inf = %10.3e' % (ex_states[2,1]+pinf_water))
```

```python
plot_riemann_solution(q_l, q_r, auxl, auxr);
```


# Beat phenomenon

```python
import vibration_toolbox as vtb

import matplotlib.pyplot as plt
from ipywidgets import interact, interactive, widgets
from IPython.display import Audio, display

# %matplotlib inline
```

In this notebook we will investigate a phenomena that occurs when the driving frequency becomes close  to  the  system’s  natural  frequency.

For zero initial conditions, the response is given by the equation:

\begin{equation}
x(t) = \frac{2f_0}{\omega^2_n - \omega^2} sin\Big(\frac{\omega_n - \omega}{2} t\Big) sin\Big(\frac{\omega_n + \omega}{2}t\Big)
\end{equation}

Since $(\omega_n - \omega)$ is small, $(\omega_n + \omega)$ is large in comparison and $sin[(\omega_n - \omega)/2]t$ oscillates with a longer period than $sin[(\omega_n + \omega)/2]t$.
The result is a rapid oscillation with slowly varying amplitude that is called a **beat**.

The system defined in the function bellow has a a natural frequency $\omega_n = 1000 \ rad/s$.

You can change the excitation frequency moving the slide 'wdr' and hear the resulting sound.

When you get close to the natural frequency you will be able to recognize the beats.

```python
@interact(wdr=widgets.FloatSlider(min=900.1, max=1100.1, step=1, value=1100))
def f(wdr):
    t, a = vtb.forced_analytical(m=1, k=10**6, x0=0, v0=0, wdr=wdr, F0=200, tf=5)
    display(Audio(data=a, rate=1/0.000125))
    plt.plot(t, a, alpha=0.5)
    plt.show()

```

Failed to display Jupyter Widget of type interactive.

  If you're reading this message in the Jupyter Notebook or JupyterLab Notebook, it may mean
  that the widgets JavaScript is still loading. If this message persists, it
  likely means that the widgets JavaScript library is either not installed or
  not enabled. See the Jupyter
  Widgets Documentation for setup instructions.

  If you're reading this message in another frontend (for example, a static
  rendering on GitHub or NBViewer),
  it may mean that your frontend doesn't currently support widgets.

```python
w = interactive(f,wdr = (900.1,1100.1,1));
display(w);
```

Failed to display Jupyter Widget of type interactive.

  If you're reading this message in the Jupyter Notebook or JupyterLab Notebook, it may mean
  that the widgets JavaScript is still loading. If this message persists, it
  likely means that the widgets JavaScript library is either not installed or
  not enabled. See the Jupyter
  Widgets Documentation for setup instructions.

  If you're reading this message in another frontend (for example, a static
  rendering on GitHub or NBViewer),
  it may mean that your frontend doesn't currently support widgets.

```python
f(wdr=10)
```

```python
w = f(wdr=10)
```

```python
@interact
def f2(wdr=10):
    t, a = vtb.forced_analytical(m=1, k=10**6, x0=0, v0=0, wdr=wdr, F0=200, tf=5)
    # display(Audio(data=a, rate=1/0.000125))
    plt.plot(t, a, alpha=0.5)
    plt.show()
    return
```

Failed to display Jupyter Widget of type interactive.

  If you're reading this message in the Jupyter Notebook or JupyterLab Notebook, it may mean
  that the widgets JavaScript is still loading. If this message persists, it
  likely means that the widgets JavaScript library is either not installed or
  not enabled. See the Jupyter
  Widgets Documentation for setup instructions.

  If you're reading this message in another frontend (for example, a static
  rendering on GitHub or NBViewer),
  it may mean that your frontend doesn't currently support widgets.

```python
w = f2(wdr=10.0)
interact(w,wdr=10.)
```

```python
interactive(f2,wdr=10.0)
```

Failed to display Jupyter Widget of type interactive.

  If you're reading this message in the Jupyter Notebook or JupyterLab Notebook, it may mean
  that the widgets JavaScript is still loading. If this message persists, it
  likely means that the widgets JavaScript library is either not installed or
  not enabled. See the Jupyter
  Widgets Documentation for setup instructions.

  If you're reading this message in another frontend (for example, a static
  rendering on GitHub or NBViewer),
  it may mean that your frontend doesn't currently support widgets.

```python

```