```python
import numpy as np
import matplotlib.pyplot as plt
from pycvodes import fpes
from pyodesys.tests._robertson import run_integration
from pyodesys.plotting import info_vlines
from pyodesys.native import native_sys
%matplotlib inline
```

```python
def integrate_and_plot(ax, title=None, info_vlines_kw=None, **kwargs):
    xout, yout, info, systems = run_integration(
        integrator='cvode', record_rhs_xvals=True, record_jac_xvals=True,
        record_order=True, record_fpe=True, **kwargs)
    info_vlines(ax, xout, info, **(info_vlines_kw or {
                'fpes': fpes, 'post_proc': (lambda x: np.exp(x)) if kwargs.get('logt', False) else None}))
    for idx, name in enumerate('ABC'):
        ax.plot(xout, yout[:, idx], label=name)
    ax.set_title((title or '') + ('%d steps, %d rhs evals, %.4f s CPU t' % (info['n_steps'], info['nfev'], info['time_cpu'])) +
                 ('' if info['success'] else ' - failed!'))
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_ylim([1e-21, 1e14])
    ax.set_xlim([1e-12, 1e19])
    ax.legend()
    return systems
```

```python
kw = dict(nsteps=18000, atol={'A': 1e-8, 'B': 1e-10, 'C': 1e-4}, rtol=1e-6,
          dep_scaling=1e4, zero_conc=1e-24, t0=1e-12, return_on_error=True)
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
integrate_and_plot(axes[0], **kw)
_ = integrate_and_plot(axes[1], first_step=1e-12, **kw)
```

```python
_ = integrate_and_plot(plt.subplot(1, 1, 1), wrapping_class=native_sys['cvode'], **kw)
```

```python
def integrate_and_plot_reduced(**kwargs):
    kw2 = kw.copy()
    kw2.update(kwargs)
    kw2['wrapping_class'] = native_sys['cvode']
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    systems = []
    for i in range(4):
        title = 'A B C'.split()
        if i > 0:
            title.pop(i - 1)
        systems.append(integrate_and_plot(
                axes.flat[i], reduced=i, title=''.join(title) + ', ', **kw2))
    return systems
```

```python
rlin = integrate_and_plot_reduced()
```

```python
kw_log = dict(atol=1e-7, logc=True, logt=True, dep_scaling=1)
rlog1 = integrate_and_plot_reduced(**kw_log)
```

```python
rloglin = integrate_and_plot_reduced(atol=1e-7, logc=True, logt=False, dep_scaling=1)
```

```python
rlinlog = integrate_and_plot_reduced(atol=1e-7, logc=False, logt=True, dep_scaling=1)
```

```python
rlog2 = integrate_and_plot_reduced(powsimp=True, **kw_log)
```

```python
import sympy as sp
sp.init_printing()
```

```python
rlog1[0][-1].exprs
```

```python
rlog1[0][-1].get_jac()
```

```python
rlog2[0][-1].exprs
```

```python
rlog2[0][-1].get_jac()
```

```python
rlog2[0][-1]._native._written_files
```

```python
def rhs_code_in_odesys(odesys):
    seen_rhs = False
    for line in open(next(filter(lambda n: n.endswith('.cpp'), odesys._native._written_files))):
        if not seen_rhs:
            if 'rhs(' in line:
                seen_rhs = True
            else:
                continue
        if 'nfev++' in line:
            yield '...'
            break
        yield line
```

```python
print(''.join(rhs_code_in_odesys(rlog1[0][-1])))
```

```python
print(''.join(rhs_code_in_odesys(rlog2[0][-1])))
```

```python

```

```python

```

```python

```


# The Riordan Group

This notebook attempt to study *The Riordan Group* article, by Shapiro, Getu, Woan and Woodson. It is my first attempt to explain ideas and to find meaning and connections among things. We reference objects and equations explained in the article using their reference numbers in order to have a direct mapping to them.

We start with some SymPy stuff in order to setup the environment:

```
import sympy
from sympy import *
from sympy.abc import x, n, z, t, y, k

from IPython.display import display, Math, Latex

init_printing() # for nice printing, a-la' TeX

%run '../riordan-group.py'
```

# Definitions

Let $M$ be an infinite matrix with elements $m_{ji} \in \mathbb{C}$ defined as:
$$ M = (m_{ji})_{i,j \geq 0}= \left[ \begin{array}{cccccc}
m_{00} & m_{01} & m_{02} & \ldots & m_{0i} & \ldots \\
m_{10} & m_{11} & m_{12} & \ldots & m_{1i} & \ldots \\
m_{20} & m_{21} & m_{22} & \ldots & m_{2i} & \ldots \\
\vdots & \vdots & \vdots & \ddots & \vdots &  \\
m_{j0} & m_{j1} & m_{j2} & \ldots & m_{ji} & \ldots \\
\vdots & \vdots & \vdots &  & \vdots & \ddots \\
\end{array} \right] $$

Now define $C_i(x)$ as a generating function taking coefficients from $i$-th column:
$$ C_i(x) = \sum_{j=0}^{\infty}{m_{ji}x^j} $$

So the following matrix product holds:
$
\left[\begin{array}{cccccc}        1 & x & x^2 & \ldots & x^j & \ldots \end{array} \right] M =
\left[\begin{array}{cccccc}        C_0(x) & C_1(x) & C_2(x) & \ldots & C_i(x)& \ldots \end{array} \right]
$

Now assume that $C_i(x)$ can be written as $g(x)f(x)^i$, for all $i \geq 0$ where:
$$ g(x) = g_0 + g_1 x + g_2 x^2 + g_3 x^3 +\ldots \\
f(x) = x + f_2 x^2 + f_3 x^3 +\ldots$$

So, there are no requirements on function $g$ while function $f$ has to obeys the following:
$$f(0) = 0  \\ \left . \frac{\partial f}{\partial x} \right|_0 = 1$$

Under these assumptions we can rewrite:
$
\left[\begin{array}{cccccc}        1 & x & x^2 & \ldots & x^j & \ldots \end{array} \right] M =
\left[\begin{array}{cccccc}        g(x) & g(x)f(x) & g(x)f(x)^2 & \ldots & g(x)f(x)^i & \ldots \end{array} \right]
$

and we call a matrix $M = (g(x), f(x))$ a *Riordan matrix*.

## Product with a vector

Now take the product on the right with a vector $\left[\begin{array}{cccccc}        a_0 & a_1 & a_2 & \ldots & a_i & \ldots \end{array} \right]^T$:
$$ \begin{split}
\left[\begin{array}{cccccc}        1 & x & x^2 & \ldots & x^j & \ldots \end{array} \right] M
\left[\begin{array}{c}        a_0 \\ a_1 \\ a_2 \\ \vdots \\ a_i \\ \vdots \end{array} \right] &=
a_0 g(x) + a_1 g(x) f(x) + a_2 g(x) f(x)^2 + \ldots \\
&= g(x) \left( a_0  + a_1 f(x) + a_2 f(x)^2 + \ldots \right) \\
&= g(x) A(f(x))
\end{split}$$
where $A(t) = a_0 + a_1 t + a_2 t^2 + \ldots$ is the formal power series of the sequence $(a_0, a_1, a_2, \ldots)$.

Observe that $ g(x) A(f(x))$ can be expanded as a formal power series $b_0 + b_1 x + b_2 x^2 + \ldots$ for some sequence $\{b_n\}_{n \in \mathbb{N}}$ after collecting coefficients for equal powers of $x$: applying $\mathcal{G}$ operator it is possibile to write $\mathcal{G}(\{b_n\}) = g(x) A(f(x))$. Since on the right side there is a product of two separable function, it does follow that the generic term of $\{b_n\}$ is a convolution by properties of $\mathcal{G}$, namely:
$$ b_n = \sum_{i=0}^{n}{g_i r_{n-i}}$$
where coefficient $r_j$ depends on both finitely coefficients $a_0, \ldots, a_j$ both on infinitely coefficients $f_2,f_3,\ldots$

## Product with a matrix

Let $M, N$ be two Riordan matrices such that $M = (g(x),f(x))$ and $N =  (h(x),l(x))$. In order to study the product $M N$ consider the $i$-th column of matrix $N$, namely coefficients in the formal power series expansion of $h(x)l(x)^i$.
By product by vector explanation above, coefficients of $i$-th column of $M N$ are denoted by series expansion of $g(x)h(f(x))l(f(x))^i$, hence repeating the same reasoning for all $i$, it is possible to write: $M N = (g(x)h(f(x)), l(f(x)))$ hence $M N$ is a Riordan matrix too.

## The set of Riordan matrices with product operation is a group

```
def g(x): return 1
def f(x): return x

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

\begin{equation} \left( 1, x \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 0 & 1 &  &  &  &  &  &  &  &  &  \\ 0 & 0 & 1 &  &  &  &  &  &  &  &  \\ 0 & 0 & 0 & 1 &  &  &  &  &  &  &  \\ 0 & 0 & 0 & 0 & 1 &  &  &  &  &  &  \\ 0 & 0 & 0 & 0 & 0 & 1 &  &  &  &  &  \\ 0 & 0 & 0 & 0 & 0 & 0 & 1 &  &  &  &  \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 &  &  &  \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 &  &  \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

## Pascal triangle standard

In this section we'll built the standard Pascal triangle, developing the Riordan matrix $(g(x),f(x))$ that represent a matrix such that row $r$ contains coefficients of $(x + y)^r$, namely ${ {r} \choose {i} }$, for $i \in \{0,\ldots,r\}$.

Since the matrix $$M = \left({ {j} \choose {i} }\right)_{i,j \geq 0}$$ is well known, we proceed in two step in order to define functions $g(x)$ and $f(x)$, starting with the former one. Observe that $g(x)$ is the same as $C_0(x) = 1 + x + x^2 + x^3 + \ldots $, so it is a formal power series, where each coefficient equals 1, hence $g(x)$ is the generating function of the sequence $\{1_n\}_{n\in \mathbb{n}}$, which has a closed form: $$g(x) = \mathcal{G}\{1_n\} = \frac{1}{1-x}$$

Now set for $f(x)$. By definition of Riordan matrix, the expansion of $C_1(x) = g(x)f(x)$ as power series should collect coefficients from the second column, namely the natural numbers. Recall that the generating function of the sequence $\{n\}_{n\in \mathbb{N}}$ is: $$\mathcal{G} \{n\} = \mathcal{G}\{n1_n\} = x\frac{\partial \mathcal{G}\{1_n\}}{\partial x} = \frac{x}{(1-x)^2}$$

Hence the following equation yields $f(x)$:
$$ \frac{1}{1-x}f(x) = \frac{x}{(1-x)^2}$$

The above reasoning is implemented with the following code:

```
def g(x): return 1/(1-x)
def f(x): return x/(1-x)

Latex(Riordan_matrix_latex_code(g, f, x, order=20))
```

```
two_var_gen_fun = 1/(1-y*(1+x))
display(series(two_var_gen_fun,(1+x)))
display(series(two_var_gen_fun,y))
for k in range(5): display(series(y**k/(1-y)**(k+1),y,n=10))
```

## Pascal triangle enhanced

In this section we approach a variant of the above triangle, starting to work with a matrix $M = (b_{nj})_{n,j \geq 0}$ where coefficients aren't explicitly defined but obey to the recurrence: $$ b_{n+1,j} = b_{n, j-1} + b_{n, j+1}$$

where $j > 0$. Assume the relation holds for all $n \in \mathbb{N}$ so it is possible to apply the *Identity principle*:
$$ \begin{split}
\mathcal{G}\{b_{n+1,j}\} &= \mathcal{G}\{b_{n, j-1} + b_{n, j+1}\} \\
\mathcal{G}\{b_{n+1,j}\} &= \mathcal{G}\{b_{n, j-1}\} + \mathcal{G}\{b_{n, j+1}\} \quad \text{ by linearity of } \mathcal{G} \\
\frac{\mathcal{G}\{b_{n,j}\}-b_{0j}}{x} &= \mathcal{G}\{b_{n, j-1}\} + \mathcal{G}\{b_{n, j+1}\} \quad \text{ by forward translation of } \mathcal{G} \\
C_j(x) &= xC_{j-1}(x) + xC_{j+1}(x) \quad  \forall j > 0: b_{0j} = 0
\end{split}$$

Another way to prove the same result, is to observe that $b_{n+1,j}$ is a *convolution* of the two terms $b_{n,j-1}$ and $b_{n,j+1}$, hence there must exists two functions, say $h_j(x) = h_{0j} + h_{1j} x + h_{2j} x^2 + \ldots$, where coefficient $h_{i,j}= b_{i,j-1} + b_{i,j+1}$ and $l(x) = l_0 + l_1 x + l_2 x^2 + \ldots$, such that:
$$\mathcal{G}\left\lbrace\sum_{k=0}^{n+1}{h_{k,j} l_{n+1-k}}\right\rbrace = h_j(x)l(x)$$

Choose $l(x) = x$ because the entire sum should reduce to $b_{n,j-1} + b_{n,j+1}$ which equals $h_{nj}$, hence:
$$\begin{split}
\mathcal{G}\left\lbrace\sum_{k=0}^{n+1}{h_{k,j} l_{n+1-k}}\right\rbrace &= \mathcal{G}\left\lbrace{h_{n,j} l_{1}}\right\rbrace \\ &=\mathcal{G}\left\lbrace{b_{n,j-1} + b_{n,j+1}}\right\rbrace \\
&=\mathcal{G}\left\lbrace{b_{n,j-1}}\right\rbrace + \mathcal{G}\left\lbrace{b_{n,j+1}}\right\rbrace \\
&= C_{j-1}(x) + C_{j+1}(x)\end{split}$$

It follows that $h_j(x) = C_{j-1}(x) + C_{j+1}(x)$ as obtained in the former derivation.

```
def g(x): return 1/sqrt(1 -4*x**2)
def f(x): return (1-sqrt(1 -4*x**2))/(2*x)

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

\begin{equation} \left( \frac{1}{\sqrt{- 4 x^{2} + 1}}, \frac{1}{2 x} \left(- \sqrt{- 4 x^{2} + 1} + 1\right) \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 0 & 1 &  &  &  &  &  &  &  &  &  \\ 2 & 0 & 1 &  &  &  &  &  &  &  &  \\ 0 & 3 & 0 & 1 &  &  &  &  &  &  &  \\ 6 & 0 & 4 & 0 & 1 &  &  &  &  &  &  \\ 0 & 10 & 0 & 5 & 0 & 1 &  &  &  &  &  \\ 20 & 0 & 15 & 0 & 6 & 0 & 1 &  &  &  &  \\ 0 & 35 & 0 & 21 & 0 & 7 & 0 & 1 &  &  &  \\ 70 & 0 & 56 & 0 & 28 & 0 & 8 & 0 & 1 &  &  \\ 0 & 126 & 0 & 84 & 0 & 36 & 0 & 9 & 0 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

```
series(f(x),x,n=20)
```

```
def g(x): return (1-x**2)/(x**2+1)
def f(x): return x/(1+x**2)

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

\begin{equation} \left( \frac{- x^{2} + 1}{x^{2} + 1}, \frac{x}{x^{2} + 1} \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 0 & 1 &  &  &  &  &  &  &  &  &  \\ -2 & 0 & 1 &  &  &  &  &  &  &  &  \\ 0 & -3 & 0 & 1 &  &  &  &  &  &  &  \\ 2 & 0 & -4 & 0 & 1 &  &  &  &  &  &  \\ 0 & 5 & 0 & -5 & 0 & 1 &  &  &  &  &  \\ -2 & 0 & 9 & 0 & -6 & 0 & 1 &  &  &  &  \\ 0 & -7 & 0 & 14 & 0 & -7 & 0 & 1 &  &  &  \\ 2 & 0 & -16 & 0 & 20 & 0 & -8 & 0 & 1 &  &  \\ 0 & 9 & 0 & -30 & 0 & 27 & 0 & -9 & 0 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

```
series(f(x),x,n=20)
```

```
def g(x): return 1/sqrt(1-2*x-3*x**2)
def f(x): return (1-x-sqrt(1-2*x-3*x**2))/(2*x)

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

\begin{equation} \left( \frac{1}{\sqrt{- \left(x + 1\right) \left(3 x - 1\right)}}, - \frac{1}{2 x} \left(x + \sqrt{- 3 x^{2} - 2 x + 1} - 1\right) \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 1 & 1 &  &  &  &  &  &  &  &  &  \\ 3 & 2 & 1 &  &  &  &  &  &  &  &  \\ 7 & 6 & 3 & 1 &  &  &  &  &  &  &  \\ 19 & 16 & 10 & 4 & 1 &  &  &  &  &  &  \\ 51 & 45 & 30 & 15 & 5 & 1 &  &  &  &  &  \\ 141 & 126 & 90 & 50 & 21 & 6 & 1 &  &  &  &  \\ 393 & 357 & 266 & 161 & 77 & 28 & 7 & 1 &  &  &  \\ 1107 & 1016 & 784 & 504 & 266 & 112 & 36 & 8 & 1 &  &  \\ 3139 & 2907 & 2304 & 1554 & 882 & 414 & 156 & 45 & 9 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

```
series(f(x),x,n=10)
display(series((3+2*x+x**2)/(1-x), n=10))
```

```
def g(x): return Integer(1)
def f(x): return x/(1-x)

for n in range(1,6):
    display(series(f(x)**n, n=10))

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

```
def d(t): return (1-t)/(1-2*t)
def h(t): return (1+t)/((1-t)*t)

Latex(Riordan_matrix_latex_code(d, h, t, order=10))
```

\begin{equation} \left( \frac{t - 1}{2 t - 1}, - \frac{t + 1}{t \left(t - 1\right)} \right) = \left[ \begin{array}{ccccccccccc}1 & 3 & 14 & 72 & 384 & 2088 & 11492 & 63780 & 356144 & 1998000 &  \\ 1 & 6 & 32 & 172 & 936 & 5144 & 28484 & 158652 & 887792 & 4986736 &  \\ 2 & 12 & 68 & 380 & 2120 & 11848 & 66388 & 372996 & 2100944 & 11860880 &  \\ 4 & 24 & 140 & 804 & 4584 & 26056 & 147956 & 840156 & 4773200 & 27138640 &  \\ 8 & 48 & 284 & 1660 & 9624 & 55512 & 319204 & 1832100 & 10504560 & 60197360 &  \\ 16 & 96 & 572 & 3380 & 19832 & 115736 & 672740 & 3899260 & 22554160 & 130268080 &  \\ 32 & 192 & 1148 & 6828 & 40392 & 237800 & 1394420 & 8150340 & 47516560 & 276462032 &  \\ 64 & 384 & 2300 & 13732 & 81672 & 483880 & 2856660 & 16812060 & 98677392 & 577867408 &  \\ 128 & 768 & 4604 & 27548 & 164408 & 978360 & 5805060 & 34348772 & 202727984 & 1193793584 &  \\ 256 & 1536 & 9212 & 55188 & 330072 & 1970040 & 11731652 & 69701820 & 413198192 & 2444317680 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

```
def curious(m): return ((1 + x)**m)*x**(m+1)/(1 -x -x**2)**(m+1)
[display(series(curious(m),n=10)) for m in range(5)]; None
```

```
def g(x): return Integer(1)
def f(x): return x*(1+x)

Latex(Riordan_matrix_latex_code(g, f, x, order=10))
```

\begin{equation} \left( 1, x \left(x + 1\right) \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 0 & 1 &  &  &  &  &  &  &  &  &  \\ 0 & 1 & 1 &  &  &  &  &  &  &  &  \\ 0 & 0 & 2 & 1 &  &  &  &  &  &  &  \\ 0 & 0 & 1 & 3 & 1 &  &  &  &  &  &  \\ 0 & 0 & 0 & 3 & 4 & 1 &  &  &  &  &  \\ 0 & 0 & 0 & 1 & 6 & 5 & 1 &  &  &  &  \\ 0 & 0 & 0 & 0 & 4 & 10 & 6 & 1 &  &  &  \\ 0 & 0 & 0 & 0 & 1 & 10 & 15 & 7 & 1 &  &  \\ 0 & 0 & 0 & 0 & 0 & 5 & 20 & 21 & 8 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

# Identities and inverse relations

```
s = series((1-x)**(2*n),x,n=10).subs(n,4)
s, s.coeff(x**5)-s.coeff(x**4)
```

```
series((1+x)**5,x,n=10)
```

```
list(map(display, (series((1/sqrt(1+2*x-3*x**2))*(1+x)**(n-1) *(1+x),x,n=10) for n in range(5))))
print("\n")
list(map(display, (series((1/sqrt(1+2*x-3*x**2))*(1+x)**(n-1),x,n=10) for n in range(5))))
None
```

---

# $A$ and $Z$ sequences for characterizing Riordan arrays

## Fibonacci triangle

```
def d(x): return 1/(1-x-x**2)
def h(x): return (1-sqrt(1-4*x))/2

Latex(Riordan_matrix_latex_code(d, h, t, order=10))
```

\begin{equation} \left( - \frac{1}{t^{2} + t - 1}, - \frac{1}{2} \left(\sqrt{- 4 t + 1} - 1\right) \right) = \left[ \begin{array}{ccccccccccc}1 &  &  &  &  &  &  &  &  &  &  \\ 1 & 1 &  &  &  &  &  &  &  &  &  \\ 2 & 2 & 1 &  &  &  &  &  &  &  &  \\ 3 & 5 & 3 & 1 &  &  &  &  &  &  &  \\ 5 & 12 & 9 & 4 & 1 &  &  &  &  &  &  \\ 8 & 31 & 26 & 14 & 5 & 1 &  &  &  &  &  \\ 13 & 85 & 77 & 46 & 20 & 6 & 1 &  &  &  &  \\ 21 & 248 & 235 & 150 & 73 & 27 & 7 & 1 &  &  &  \\ 34 & 762 & 741 & 493 & 258 & 108 & 35 & 8 & 1 &  &  \\ 55 & 2440 & 2406 & 1644 & 903 & 410 & 152 & 44 & 9 & 1 &  \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \ddots\end{array} \right] \end{equation}

```
Eq(h(t).factor(), series(h(t),n=10))
```

```
one_less_two_powers = x/((1-x)*(1-2*x))
taken_apart = one_less_two_powers.apart()

display(taken_apart)
display(series(1/(1-2*x)))
display(series(1/(x-1)))

series(one_less_two_powers, n=10)
```

```
def autocorrelation_polynomial(x=x,y=y): return (1 + x**2 * y**2 + x**3 * y**2)
def avoiding_pattern_expr(x=x,y=y): return autocorrelation_polynomial(x,y) / ((1-x-y)*autocorrelation_polynomial(x,y) + x**4 * y**2)

display(avoiding_pattern_expr())
display(avoiding_pattern_expr(x,t/x))

for_thm = avoiding_pattern_expr(x,t)
for_thm.series(x,n=2).coeff(x).series()
```

```

# getting the series over *x* means asking for *rows*'s generating functions
x_series = series(avoiding_pattern_expr(), x,n=10)

# get generating function of row 4th, expand it 'column-wise' (over *y*) in order to build row 4th as desired
avoiding_pattern_expr(), x_series.coeff(x**4).series(y,n=10)
```

```
series(1/(1-t),t,n=10)
```

```
# pattern p: 011
def autocorrelation_polynomial(x=x,y=y): return 1
def avoiding_pattern_expr(x=x,y=y): return autocorrelation_polynomial(x,y) / ((1-x-y)*autocorrelation_polynomial(x,y) + x**2 * y**1)

display(avoiding_pattern_expr())
display(avoiding_pattern_expr(x,t/x))

for_thm = avoiding_pattern_expr(x,t/x)
for_thm.series(x,n=2),for_thm.series(t,n=2)
```

```
def R(n,k): return k/(2*(k-n-1))
[display(series(R(n,k),x=n,n=10)) for k in range(10)]
[display(series(R(n,k),x=k,n=10)) for n in range(10)]

display(R(n, k))
```

```
def R(n,k): return (k-1)/n
[display(series(R(n,k),x=n,n=10)) for k in range(10)]
[display(series(R(n,k),x=k,n=10)) for n in range(10)]

display(R(n, k))
```

```
def f_1(n): return factorial(3*n+1)*factorial(2*n-5)*factorial(n+2)**2
simplify(Sum(f_1(n-k)/f_1(n), (k,0,3)).doit())
```

---

# Printing and displaying stuff

In this foot section we save some display tricks in order to support the implementation of a function that builds Riordan coefficient matrix:

```
from IPython.display import Math
Math(r'F(k) = \int_{-\infty}^{\infty} f(x) e^{2\pi i k} dx')
```

$$F(k) = \int_{-\infty}^{\infty} f(x) e^{2\pi i k} dx$$

```
from IPython.display import Latex
Latex(r"""\begin{eqnarray}
\nabla \times \vec{\mathbf{B}} -\, \frac1c\, \frac{\partial\vec{\mathbf{E}}}{\partial t} & = \frac{4\pi}{c}\vec{\mathbf{j}} \\
\nabla \cdot \vec{\mathbf{E}} & = 4 \pi \rho \\
\nabla \times \vec{\mathbf{E}}\, +\, \frac1c\, \frac{\partial\vec{\mathbf{B}}}{\partial t} & = \vec{\mathbf{0}} \\
\nabla \cdot \vec{\mathbf{B}} & = 0
\end{eqnarray}""")
Latex(r"""\[ 4 \]""")
```

\[ 4 \]

```
from IPython.display import display, Math, Latex
display(Math(r'F(k) = \int_{-\infty}^{\infty} f(x) e^{2\pi i k} dx'))
```

$$F(k) = \int_{-\infty}^{\infty} f(x) e^{2\pi i k} dx$$


# Implementation of a Devito viscoacoustic equations

## This tutorial is contributed by SENAI CIMATEC (2020)

This tutorial is based on:

**Linear inversion in layered viscoacoustic media using a time‐domain method** (1994)

Joakim O. Blanch and William W. Symes

SEG Technical Program Expanded Abstracts

https://doi.org/10.1190/1.1822695

**True-amplitude prestack depth migration** (2007)

Feng Deng and George A. McMechan

GEOPHYSICS Technical Papers

https://doi.org/10.1190/1.2714334

**Attenuation compensation for least-squares reverse time migration using the viscoacoustic-wave equation** (2014)

Gaurav Dutta and Gerard T. Schuster

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2013-0414.1

**Multiscale viscoacoustic waveform inversion with the second generation wavelet transform and adaptive time–space domain finite-difference method** (2014)

Zhiming Ren, Yang Liu,and Qunshan Zhang

Geophysical Journal International, Volume 197, Issue 2, 1 May 2014, Pages 948–974

https://doi.org/10.1093/gji/ggu024

**Viscoacoustic prestack reverse time migration based on the optimal time-space domain high-order finite-difference method** (2014)

Yan Zhao, Yang Liu, and Zhi-Ming Ren

Appl. Geophys. 11, 50–62.

https://doi.org/10.1007/s11770-014-0414-8

**A stable and efficient approach of Q reverse time migration** (2018)

Yan Zhao, Ningbo Mao, and Zhiming Ren

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2018-0022.1

## Introduction

The conversion of mechanical energy to heat, occurs during the propagation of seismic waves on the subsurface, due to the viscosity of the rocks. The presence of oil and gas in these rocks causes seismic attenuations. Thus, associated effects, such as dispersion and dissipation, can significantly affect the amplitudes, as well as the phase of the seismic pulse.  However, in the seismic exploration, the subsurface has still been considered as an ideal elastic/acoustic medium, that is, disregarding its mitigating effect. In practice, the propagation of seismic waves on the subsurface is in many ways different from propagation in an ideal solid.

For example, some subsurface rocks have anisotropic properties, are heterogeneous, porous and so on. The acoustic/elastic wave equation is not sensitive enough to describe propagation in these more complicated mediums. Generally, the viscosity of materials in the subsurface causes energy dissipation and consequently a decrease in amplitude, in addition to modifying the frequency content of the waves. This phenomenon of energy dissipation of the wave is called seismic absorption or attenuation.

The goal of this tutorial is to perform a seismic modeling taking into account the viscosity of the medium, so that it is possible to more accurately simulate the seismic data and consequently build images with better resolution in the processing of this data, in addition to extracting more detailed information on rocky materials through seismic inversion.

This tutorial follow three main viscoacoustic approaches in time-space domain:

- Blanch and Symes (1995) / Dutta and Schuster (2014)

- Ren et al. (2014)

- Deng and McMechan (2007)

Table of symbols

| Symbol         | Description
| :---         | :---
|$f$                                  |Frequency                                            |
|$f_o$                                |Reference frequency                                  |
|$\omega$                             |Angular frenquency                                   |
|$\omega_0$                           |Angular Reference Frequency                          |
|$v$                                  |Velocity model                                       |
|$v_0$                                |Reference velocity at $\omega_0$                     |
|$\kappa$                             |Bulk modulus                                         |
|$g$                                  |Absorption coefficient                               |
|$\tau$                               |Relaxation time                                      |
|$\tau_\sigma$                        |Stress relaxation parameter                          |
|$\tau_\epsilon$                      |Strain relaxation parameter                          |
|$Q$                                  |Quality factor                                       |
|$\eta$                               |Viscosity                                            |
|$\rho$                               |Density                                              |
|$\nabla$                             |Nabla operator                                       |
|$P({\bf x},t)$                       |Pressure field                                       |
|$r({\bf x},t)$                       |Memory variable                                      |
|${\bf v}({\bf x},t)$                 |Particle velocity                                    |
|$S({\bf x}_s,t)$                     |Source                                               |

# Seismic modelling with Devito

Before start with the viscoacoustic approaches we will describe a setup of seismic modelling with Devito in a simple 2D case. We will create a physical model of our domain and define a single source and an according set of receivers to model for the forward model. But first, we initialize some basic utilities.

```python
import numpy as np
import sympy as sp

import matplotlib.pyplot as plt

from devito import *
from examples.seismic.source import RickerSource, WaveletSource, TimeAxis
from examples.seismic import ModelViscoacoustic, plot_image, setup_geometry, plot_velocity
```

```python
nx = 300
nz = 300

# Define a physical size
shape = (nx, nz)
spacing = (20., 20.)
origin = (0., 0.)
nlayers = 3
nbl = 50
space_order = 8
dtype = np.float32

# Model physical parameters:
vp = np.zeros(shape)
qp = np.zeros(shape)
rho = np.zeros(shape)

# Define a velocity profile. The velocity is in km/s
vp_top = 1.5
vp_bottom = 3.5

# Define a velocity profile in km/s
v = np.empty(shape, dtype=dtype)
v[:] = vp_top  # Top velocity (background)
vp_i = np.linspace(vp_top, vp_bottom, nlayers)
for i in range(1, nlayers):
    v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity

qp[:] = 3.516*((v[:]*1000.)**2.2)*10**(-6) # Li's empirical formula

rho[:] = 0.31*(v[:]*1000.)**0.25 # Gardner's relation
```

```python
#NBVAL_IGNORE_OUTPUT
model = ModelViscoacoustic(space_order=space_order, vp=v, qp=qp, b=1/rho,
                           origin=origin, shape=shape, spacing=spacing,
                           nbl=nbl)
```

    Operator `initdamp` ran in 0.01 s
    Operator `pad_vp` ran in 0.01 s
    Operator `pad_b` ran in 0.02 s
    Operator `pad_qp` ran in 0.01 s

```python
#NBVAL_IGNORE_OUTPUT
aspect_ratio = model.shape[0]/model.shape[1]

plt_options_model = {'cmap': 'jet', 'extent': [model.origin[0], model.origin[0] + model.domain_size[0],
                                               model.origin[1] + model.domain_size[1], model.origin[1]]}
fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]

img1 = ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, **plt_options_model)
fig.colorbar(img1, ax=ax[0])
ax[0].set_title(r"V (km/s)", fontsize=20)
ax[0].set_xlabel('X (m)', fontsize=20)
ax[0].set_ylabel('Depth (m)', fontsize=20)
ax[0].set_aspect('auto')

img2 = ax[1].imshow(np.transpose(qp), vmin=15, vmax=220, **plt_options_model)
fig.colorbar(img2, ax=ax[1])
ax[1].set_title("Q", fontsize=20)
ax[1].set_xlabel('X (m)', fontsize=20)
ax[1].set_ylabel('Depth (m)', fontsize=20)
ax[1].set_aspect('auto')

img3 = ax[2].imshow(np.transpose(rho), vmin=1.9, vmax=2.4, **plt_options_model)
fig.colorbar(img3, ax=ax[2])
ax[2].set_title(r"Density $\rho$ (g/cm^3)", fontsize=20)
ax[2].set_xlabel('X (m)', fontsize=20)
ax[2].set_ylabel('Depth (m)', fontsize=20)
ax[2].set_aspect('auto')

plt.tight_layout()
```

```python
f0 = 0.005 # peak/dominant frequency
b = model.b
rho = 1./b

# velocity model
vp = model.vp
lam = vp * vp * rho

t_s = (sp.sqrt(1.+1./model.qp**2)-1./model.qp)/f0
t_ep = 1./(f0**2*t_s)
tt = (t_ep/t_s) - 1.

s = model.grid.stepping_dim.spacing
damp = model.damp
```

```python
# Time step in ms and time range:
t0, tn = 0., 2000.
dt = model.critical_dt
time_range = TimeAxis(start=t0, stop=tn, step=dt)
```

```python
from examples.seismic import Receiver

def src_rec(p, model):
    src = RickerSource(name='src', grid=model.grid, f0=f0, time_range=time_range)
    src.coordinates.data[0, :] = np.array(model.domain_size) * .5
    src.coordinates.data[0, -1] = 8.

    # Create symbol for receivers
    rec = Receiver(name='rec', grid=model.grid, npoint=shape[0], time_range=time_range)

    # Prescribe even spacing for receivers along the x-axis
    rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=shape[0])
    rec.coordinates.data[:, 1] = 8.

    src_term = src.inject(field=p.forward, expr=(s*src))
    rec_term = rec.interpolate(expr=p)

    return src_term + rec_term, src, rec
```

Auxiliary functions for plotting data:

```python
def plot_receiver(rec):
    rec_plot = rec.resample(num=1001)
    scale_for_plot = np.diag(np.linspace(1.0, 2.5, 1001)**2.0)
    # Pressure (txx + tzz) data at sea surface
    extent = [rec_plot.coordinates.data[0, 0], rec_plot.coordinates.data[-1, 0], 1e-3*tn, t0]
    aspect = rec_plot.coordinates.data[-1, 0]/(1e-3*tn)/.5

    plt.figure(figsize=(10, 10))
    plt.imshow(np.dot(scale_for_plot, rec_plot.data), vmin=-.01, vmax=.01, cmap="seismic",
           interpolation='lanczos', extent=extent, aspect=aspect)
    plt.ylabel("Time (s)", fontsize=20)
    plt.xlabel("Receiver position (m)", fontsize=20)
```

```python
def plot_v_and_p(model, v, p):

    slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]
    scale = .5*1e-3

    plt_options_model = {'extent': [model.origin[0] , model.origin[0] + model.domain_size[0],
                                    model.origin[1] + model.domain_size[1], model.origin[1]]}

    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 7))

    ax[0].imshow(np.transpose(v[0].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[0].set_aspect('auto')
    ax[0].set_xlabel('X (m)', fontsize=20)
    ax[0].set_ylabel('Depth (m)', fontsize=20)
    ax[0].set_title(r"$v_{x}$", fontsize=20)

    ax[1].imshow(np.transpose(v[1].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[1].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[1].set_aspect('auto')
    ax[1].set_xlabel('X (m)', fontsize=20)
    ax[1].set_title(r"$v_{z}$", fontsize=20)

    ax[2].imshow(np.transpose(p.data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[2].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[2].set_aspect('auto')
    ax[2].set_xlabel('X (m)', fontsize=20)
    ax[2].set_title(r"$P$", fontsize=20)
```

## Equation based on standard linear solid (SLS) rheological model

The equations of motion for a viscoacoustic medium can be written as:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\tau + 1)(\nabla \cdot {\bf v}) + r = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\frac{\partial r}{\partial t} + \frac{1}{\tau_{\sigma}} [r + \tau \kappa (\nabla \cdot {\bf v})] = 0.
\end{array}
\right.
\end{equation}

Where $\tau = \tau_{\epsilon}/\tau_{\sigma} -1$ represents the magnitude of  $Q$. $\tau_{\epsilon}$ and $\tau_{\sigma}$ are, respectively, the stress and strain relaxation parameters, given by:

\begin{equation}
\tau_\sigma = \frac{\sqrt{Q^2+1}-1}{2 \pi f_0 Q}
\end{equation}
and
\begin{equation}
\tau_\epsilon= \frac{\sqrt{Q^2+1}+1}{2\pi f_0 Q}
\end{equation}

```python
# Stencil created from Blanch and Symes (1995) / Dutta and Schuster (2014)
def SLS(model, p, r, v):

    # Bulk modulus
    bm = rho * (vp * vp)

    # Define PDE to v
    pde_v = v.dt + b * grad(p)
    u_v = Eq(v.forward, damp * solve(pde_v, v.forward))

    # Define PDE to r
    pde_r = r.dt + (1. / t_s) * (r + tt * bm * div(v.forward))
    u_r = Eq(r.forward, damp * solve(pde_r, r.forward))

    # Define PDE to p
    pde_p = p.dt + bm * (tt + 1.) * div(v.forward) + r.forward
    u_p = Eq(p.forward, damp * solve(pde_p, p.forward))

    return [u_v, u_r, u_p]
```

```python
# Seismic Modelling from Blanch and Symes (1995) / Dutta and Schuster (2014) viscoacoustic wave equation.
def modelling_SLS(model):

    # Create symbols for particle velocity, pressure field, memory variable, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    r = TimeFunction(name="r", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = SLS(model, p, r, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_SLS(model)
```

    Operator `Kernel` ran in 2.72 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.87797, atol=1e-3, rtol=0)
```

## Equation based on Kelvin-Voigt (KV) rheological model

The viscoacoustic wave equation in time domain is written as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - v^{2}\nabla^{2}{P} - \eta\nabla^{2}\left(\frac{\partial P}{\partial t}\right) = S({\bf x}_{s}, t),
\end{equation}

where $\eta = \frac{v^2}{\omega_{0}Q}$ represents the viscosity of medium.

Considering the variable density $\rho$, the equation can be rewritten as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - \kappa \nabla \cdot \frac{1}{\rho} \nabla{P} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla \left(\frac{\partial{P}}{\partial{t}}\right) = S({\bf x}_{s}, t).
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{ll}
\frac{\partial P}{\partial t} + \kappa \nabla \cdot {\bf v} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla{P} = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho} \nabla{P} = 0
\end{array}
\right.
\end{equation}

```python
# Stencil created from Ren et al. (2014) viscoacoustic wave equation.
def KV(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v.dt + b * grad(p)
    u_v = Eq(v.forward, damp * solve(pde_v, v.forward))

    # Define PDE to p
    pde_p = p.dt + lam * div(v.forward) - (lam / (w * model.qp)) * div(b * grad(p, shift=.5), shift=-.5)

    u_p = Eq(p.forward, damp * solve(pde_p, p.forward))

    return [u_v, u_p]
```

```python
# Seismic Modelling from Ren et al. (2014) viscoacoustic wave equation.
def modelling_KV(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = KV(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_KV(model)
```

    Operator `Kernel` ran in 2.02 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 15, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.0639238, atol=1e-3, rtol=0)
```

## Equation based on Maxwell rheological model

The viscoacoustic wave equation for the propagating pressure $P$ in the time-space domain:

\begin{equation}
\frac{1}{v^2}\frac{\partial^{2}P}{\partial{t^2}} - \nabla^{2}P + \frac{g}{v}\frac{\partial P}{\partial{t}}  = S({\bf x}_{s}, t),
\end{equation}

where $g$ is the absorption coefficient, given by:

\begin{equation}
g = \frac{2\pi f_{0}}{vQ},
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\nabla \cdot {\bf v}) + \frac{2\pi f_{0}}{Q}P= S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\end{array}
\right.
\end{equation}

```python
# Stencil created from Deng and McMechan (2007) viscoacoustic wave equation.
def Maxwell(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v.dt + b * grad(p)
    u_v = Eq(v.forward, damp * solve(pde_v, v.forward))

    # Define PDE to p
    pde_p = p.dt + lam * div(v.forward) + (w / model.qp) * p
    u_p = Eq(p.forward, damp * solve(pde_p, p.forward))

    return [u_v, u_p]
```

```python
# Seismic Modelling from Deng and McMechan (2007) viscoacoustic wave equation.
def modelling_Maxwell(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = Maxwell(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_Maxwell(model)
```

    Operator `Kernel` ran in 0.77 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.1323929, atol=1e-3, rtol=0)
```

# More references

[1] https://academic.oup.com/gji/article/197/2/948/616510

[2] https://link.springer.com/article/10.1007/s11770-014-0414-8

[3] https://janth.home.xs4all.nl/Software/fdelmodcManual.pdf


# Self-Driving Car Engineer Nanodegree

## Project 2: **Advanced Lane Finding**
***
The goals / steps of this project are the following:

* Compute the camera calibration matrix and distortion coefficients given a set of chessboard images.
* Apply a distortion correction to raw images.
* Use color transforms, gradients, etc., to create a thresholded binary image.
* Apply a perspective transform to rectify binary image ("birds-eye view").
* Detect lane pixels and fit to find the lane boundary.
* Determine the curvature of the lane and vehicle position with respect to center.
* Warp the detected lane boundaries back onto the original image.
* Output visual display of the lane boundaries and numerical estimation of lane curvature and vehicle position.

## Import Packages

```python
#importing some useful packages
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
import cv2
import os
%matplotlib inline
```

## Read in a Sample Image

```python
#reading in an image
image = mpimg.imread('test_images/straight_lines1.jpg')

#printing out some stats and plotting
print('This image is:', type(image), 'with dimensions:', image.shape)
plt.imshow(image)  # if you wanted to show a single color channel image called 'gray', for example, call as plt.imshow(gray, cmap='gray')
```

## Helper Functions from Project 1

```python
import math

def grayscale(img):
    """
    Applies the Grayscale transform
    This will return an image with only one color channel
    but NOTE: to see the returned image as grayscale
    (assuming your grayscaled image is called 'gray')
    you should call plt.imshow(gray, cmap='gray')
    """
    return cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    # Or use BGR2GRAY if you read an image with cv2.imread()
    # return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

def region_of_interest(img, vertices):
    """
    Applies an image mask.

    Only keeps the region of the image defined by the polygon
    formed from `vertices`. The rest of the image is set to black.
    `vertices` should be a numpy array of integer points.
    """
    #defining a blank mask to start with
    mask = np.zeros_like(img)

    #defining a 3 channel or 1 channel color to fill the mask with depending on the input image
    if len(img.shape) > 2:
        channel_count = img.shape[2]  # i.e. 3 or 4 depending on your image
        ignore_mask_color = (255,) * channel_count
    else:
        ignore_mask_color = 255

    #filling pixels inside the polygon defined by "vertices" with the fill color
    cv2.fillPoly(mask, vertices, ignore_mask_color)

    #returning the image only where mask pixels are nonzero
    masked_image = cv2.bitwise_and(img, mask)
    masked_color = np.dstack((masked_image, masked_image, masked_image)) * 255

    return masked_image, masked_color

def weighted_img(img, initial_img, α=0.8, β=1., γ=0.):
    """
    `img` is the output of the hough_lines(), An image with lines drawn on it.
    Should be a blank image (all black) with lines drawn on it.

    `initial_img` should be the image before any processing.

    The result image is computed as follows:

    initial_img * α + img * β + γ
    NOTE: initial_img and img must be the same shape!
    """
    return cv2.addWeighted(initial_img, α, img, β, γ)
```

## Camera Calibration
### Extract Chessboard Corners

```python
import glob

# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
nrows = 6
ncols = 9
objp = np.zeros((nrows*ncols,3), np.float32)
objp[:,:2] = np.mgrid[0:ncols, 0:nrows].T.reshape(-1,2)

# Arrays to store object points and image points from all the images.
objpoints = [] # 3d points in real world space
imgpoints = [] # 2d points in image plane.

# Make a list of calibration images
images = glob.glob('camera_cal/calibration*.jpg')

# Step through the list and search for chessboard corners
for idx, fname in enumerate(images):
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Find the chessboard corners
    ret, corners = cv2.findChessboardCorners(gray, (ncols,nrows), None)

    # If found, add object points, image points
    if ret == True:
        objpoints.append(objp)
        imgpoints.append(corners)

        # Draw and display the corners
        cv2.drawChessboardCorners(img, (ncols,nrows), corners, ret)
        #write_name = 'corners_found'+str(idx)+'.jpg'
        #cv2.imwrite(write_name, img)
        #
        #Uncomment the two lines below to view the result
        #cv2.imshow('img', img)
        #cv2.waitKey(500)

cv2.destroyAllWindows()
```

### Compute Camera Calibration Matrix

```python
import pickle
%matplotlib inline

# Test undistortion on an image
img = cv2.imread('camera_cal/calibration1.jpg')
img_size = (img.shape[1], img.shape[0])

# Do camera calibration given object points and image points
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, img_size, None, None)

dst = cv2.undistort(img, mtx, dist, None, mtx)
cv2.imwrite('output_images/undistorted_chessboard.jpg',dst)

# Save the camera calibration result for later use
dist_pickle = {}
dist_pickle["mtx"] = mtx
dist_pickle["dist"] = dist
pickle.dump( dist_pickle, open( "camera_cal/wide_dist_pickle.p", "wb" ) )
#dst = cv2.cvtColor(dst, cv2.COLOR_BGR2RGB)
# Visualize undistortion
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(20,10))
ax1.imshow(img)
ax1.set_title('Original Image', fontsize=30)
ax2.imshow(dst)
ax2.set_title('Undistorted Image', fontsize=30)
```

### Derivation of polynomial transformation from image coordinates to real coordinates

While most of the content in the helper functions is based on code and theory described in the lessons,
it might be useful to describe the logic used to tranform polynomials from the image coordinates to the
real coordinates since this was not explicitly described in the course lessons. This section therefore,
provides a brief description of the approach used to perform the tranformation.

Let the equation of the polynomial in the image coordinates be:

\begin{equation}
x_i = A_i y_i^{2} + B_i y_i + C_i
\end{equation}

In the above equation, $x$ is the dependent variable and $y$ is the independent
variable. Let the equation of the corresponding polynomial in the real coordinates be:

\begin{equation}
x_r = A_r y_r^{2} + B_r y_r + C_r
\end{equation}

Let $\lambda_x$ and $\lambda_y$ be the $x$ and $y$ distances (in meters) per pixel. The $x$ and $y$
coordinates can therefore be transformed as:

\begin{equation}
x_r = \lambda_{x} x_i  \\
y_r = \lambda_{y} x_i
\end{equation}

If we substitute the above equations in to the polynomial expression in
real coordinates, we get:

\begin{equation}
x_i \lambda_x = A_i\lambda_y^2 y_i^{2} + B_i \lambda_y y_i + C_i \\
x_i = \frac{A_i\lambda_y^{2}}{\lambda_x} y_i^{2} + \frac{B_i \lambda_y}{\lambda_x} y_i + \frac{C_i}{\lambda_x}
\end{equation}

Comparing the above equation with the equation for the image coordinates and equating the coefficients, we get:

\begin{equation}
A_r = \frac{A_i\lambda_y^{2}}{\lambda_x} \\
B_r = \frac{B_i \lambda_y}{\lambda_x} \\
C_r = \frac{C_i}{\lambda_x}
\end{equation}

### Helper Functions for Project 2

```python
def edge_thresholds(img):
    # Extract R-channel
    R = img[:,:,0]

    # Extract S-channel
    hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
    s_channel = hls[:,:,2]

    # Grayscale image
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    # Sobel x
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0) # Take the derivative in x
    abs_sobelx = np.absolute(sobelx) # Absolute x derivative to accentuate lines away from horizontal
    scaled_sobel = np.uint8(255*abs_sobelx/np.max(abs_sobelx))

    # Threshold x gradient
    thresh_min = 50
    thresh_max = 100
    sxbinary = np.zeros_like(scaled_sobel)
    sxbinary[(scaled_sobel >= thresh_min) & (scaled_sobel <= thresh_max)] = 1

    # Threshold R channel
    thresh_R = (150, 255)
    binary_R = np.zeros_like(R)
    binary_R[(R > thresh_R[0]) & (R <= thresh_R[1])] = 1

    # Threshold S channel
    s_thresh_min = 100
    s_thresh_max = 255
    s_binary = np.zeros_like(s_channel)
    s_binary[(s_channel >= s_thresh_min) & (s_channel <= s_thresh_max)] = 1

    # Combine multiple binary images
    combined_binary = np.zeros_like(sxbinary)
    combined_binary[((s_binary == 1) & binary_R == 1) | (sxbinary == 1) ] = 1
    #combined_binary[((s_binary == 1) & binary_R == 1)] = 1
    #combined_binary[((s_binary == 1)) | (sxbinary == 1) ] = 1

    color_image = np.dstack((combined_binary, combined_binary, combined_binary)) * 255
    return binary_R, s_binary, sxbinary, combined_binary, color_image

def topview(image):
    """
    Perform perspective transformation to obtain a top view of the image
    """

    src = np.float32([[600, 444], [675, 444], [1041, 676], [268, 676]])
    offsetv = 0
    offseth = 300
    img_size = (image.shape[1], image.shape[0])
    dst = np.float32([[offseth, offsetv], [img_size[0]-offseth, offsetv],
                                     [img_size[0]-offseth, img_size[1]-offsetv],
                                     [offseth, img_size[1]-offsetv]])

    M = cv2.getPerspectiveTransform(src, dst)
    Minv = cv2.getPerspectiveTransform(dst, src)
    warped = cv2.warpPerspective(image, M, image.shape[1::-1], flags=cv2.INTER_LINEAR)

    return warped, M, Minv

def fit_poly(img_shape, leftx, lefty, rightx, righty):

    # Fit a second order polynomial to each with np.polyfit() ###
    left_fit = np.polyfit(lefty, leftx, 2)
    right_fit = np.polyfit(righty, rightx, 2)

    # Generate x and y values for plotting
    ploty = np.linspace(0, img_shape[0]-1, img_shape[0])

    #Calc both polynomials using ploty, left_fit and right_fit ###
    left_fitx = left_fit[0]*ploty**2 + left_fit[1]*ploty + left_fit[2]
    right_fitx = right_fit[0]*ploty**2 + right_fit[1]*ploty + right_fit[2]
    middle_fitx = 0.5*(left_fitx + right_fitx)
    middle_fit = np.polyfit(middle_fitx, ploty, 2)

    return left_fitx, right_fitx, middle_fitx, ploty, left_fit, right_fit, middle_fit

def find_lane_pixels(binary_warped):
    # Take a histogram of the bottom half of the image
    histogram = np.sum(binary_warped[binary_warped.shape[0]//2:,:], axis=0)

    # Create an output image to draw on and visualize the result
    stacked_windows_image = np.dstack((binary_warped, binary_warped, binary_warped))

    # Find the peak of the left and right halves of the histogram
    # These will be the starting point for the left and right lines
    midpoint = np.int(histogram.shape[0]//2)
    leftx_base = np.argmax(histogram[:midpoint])
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # Choose the number of sliding windows
    nwindows = 9
    # Set the width of the windows +/- margin
    margin = 100
    # Set minimum number of pixels found to recenter window
    minpix = 50

    # Set height of windows - based on nwindows above and image shape
    window_height = np.int(binary_warped.shape[0]//nwindows)

    # Identify the x and y positions of all nonzero pixels in the image
    nonzero = binary_warped.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])

    # Current positions to be updated later for each window in nwindows
    leftx_current = leftx_base
    rightx_current = rightx_base

    # Create empty lists to receive left and right lane pixel indices
    left_lane_inds = []
    right_lane_inds = []

    # Step through the windows one by one
    for window in range(nwindows):
        # Identify window boundaries in x and y (and right and left)
        win_y_low = binary_warped.shape[0] - (window+1)*window_height
        win_y_high = binary_warped.shape[0] - window*window_height

        win_xleft_low = leftx_current - margin
        win_xleft_high = leftx_current + margin
        win_xright_low = rightx_current - margin
        win_xright_high = rightx_current + margin

        # Draw the windows on the visualization image
        cv2.rectangle(stacked_windows_image,(win_xleft_low,win_y_low),
                              (win_xleft_high,win_y_high),(0,255,0), 2)
        cv2.rectangle(stacked_windows_image,(win_xright_low,win_y_low),
                              (win_xright_high,win_y_high),(0,255,0), 2)

        good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
        (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high)).nonzero()[0]
        good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
        (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high)).nonzero()[0]

        # Append these indices to the lists
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)

        #If number of pixes in windows exceeds minpix pixels, recenter next window
        if len(good_left_inds) > minpix:
            leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
        if len(good_right_inds) > minpix:
            rightx_current = np.int(np.mean(nonzerox[good_right_inds]))

    # Concatenate the arrays of indices (previously was a list of lists of pixels)
    try:
        left_lane_inds = np.concatenate(left_lane_inds)
        right_lane_inds = np.concatenate(right_lane_inds)
    except ValueError:
        # Avoids an error if the above is not implemented fully
        pass

    stacked_windows_image[nonzeroy[left_lane_inds], nonzerox[left_lane_inds]] = [255, 0, 0]
    stacked_windows_image[nonzeroy[right_lane_inds], nonzerox[right_lane_inds]] = [0, 0, 255]

    # Extract left and right line pixel positions
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds]
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    return leftx, lefty, rightx, righty, left_lane_inds, right_lane_inds, stacked_windows_image

def search_around_poly(binary_warped):
    # Grab activated pixels
    nonzero = binary_warped.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])

    # Extract the lane points
    leftx, lefty, rightx, righty, left_lane_inds, right_lane_inds, stacked_windows_image = find_lane_pixels(binary_warped)

    # Fit new polynomials
    left_fitx, right_fitx, middle_fitx, ploty, left_fit, right_fit, middle_fit = fit_poly(binary_warped.shape, leftx, lefty, rightx, righty)

    # Create an image to draw on and an image to show the selection window
    binary_warped_color = np.dstack((binary_warped, binary_warped, binary_warped))*255
    window_img = np.zeros_like(binary_warped_color)

    # Identify the region between the left and right lanes
    left_region_limits = np.array([np.transpose(np.vstack([left_fitx, ploty]))])
    right_region_limits = np.array([np.flipud(np.transpose(np.vstack([right_fitx, ploty])))])
    region_pts = np.hstack((left_region_limits, right_region_limits))

    # Draw the polynomial fit for middle lane line
    mid_lower_limits = np.array([np.transpose(np.vstack([middle_fitx-5, ploty]))])
    mid_upper_limits = np.array([np.flipud(np.transpose(np.vstack([middle_fitx+5, ploty])))])
    mid_region_pts = np.hstack((mid_lower_limits, mid_upper_limits))

    # Draw the polynomial fit for left lane line
    left_lower_limits = np.array([np.transpose(np.vstack([left_fitx-5, ploty]))])
    left_upper_limits = np.array([np.flipud(np.transpose(np.vstack([left_fitx+5, ploty])))])
    left_region_pts = np.hstack((left_lower_limits, left_upper_limits))

    # Draw the polynomial fit for right lane line
    right_lower_limits = np.array([np.transpose(np.vstack([right_fitx-5, ploty]))])
    right_upper_limits = np.array([np.flipud(np.transpose(np.vstack([right_fitx+5, ploty])))])
    right_region_pts = np.hstack((right_lower_limits, right_upper_limits))

    # Write left, middle and right polynomial fits to image using different colors
    cv2.fillPoly(stacked_windows_image, np.int_([mid_region_pts]), (255,255, 0)) #yellow
    cv2.fillPoly(stacked_windows_image, np.int_([left_region_pts]), (255, 105, 180)) #pink
    cv2.fillPoly(stacked_windows_image, np.int_([right_region_pts]), (128, 0, 128)) #purple

    # Draw the lane onto the warped blank image
    cv2.fillPoly(window_img, np.int_([region_pts]), (0,255, 0))
    result = cv2.addWeighted(binary_warped_color, 1, window_img, 0.3, 0)

    # Draw the lane onto the warped blank image
    result[nonzeroy[left_lane_inds], nonzerox[left_lane_inds]] = [255, 0, 0]
    result[nonzeroy[right_lane_inds], nonzerox[right_lane_inds]] = [0, 0, 255]

    return result, stacked_windows_image, left_fit, right_fit, middle_fit, \
                                          left_fitx, right_fitx, middle_fitx, ploty

def convert_to_real(poly_fit, xm_per_pix, ym_per_pix, poly_degree):
    """
    Convert from image coordinates to real coordinates using scaling factors
    """
    poly_fit_real = [0.0] * (poly_degree+1)

    for d in range(poly_degree+1):
        poly_fit_real[d] = poly_fit[d] * xm_per_pix / (ym_per_pix ** (poly_degree - d))

        return poly_fit_real

def measure_curvature(poly_fit, y_eval):
    '''
    Calculates the curvature of polynomial given the polynomial coefficients
    '''
    poly_curverad = ((1.0 + (2.0*poly_fit[0]*y_eval + poly_fit[1])**2)**1.5)/(2.0*poly_fit[0])

    return poly_curverad

def process_single_image(image, write_to_file = False):
    alpha = 1
    beta = 1
    gamma = 0

    # Specify corners of the quadrilateral masked region of interest
    imshape = image.shape
    mask_vertices = np.array([[(0.05*imshape[1],imshape[0]),
                      (0.47*imshape[1], 0.6*imshape[0]),
                      (0.53*imshape[1], 0.6*imshape[0]),
                      (0.95*imshape[1], imshape[0])]], dtype=np.int32)

    # Pixel to real world conversion factors
    ym_per_pix = 30/720 # meters per pixel in y dimension
    xm_per_pix = 3.7/700 # meters per pixel in x dimension

    # Write the original image to disk for reference
    if write_to_file:
        cv2.imwrite('output_images/original_image.png', cv2.cvtColor(image, cv2.COLOR_RGB2BGR))

    # Undistort the original image by applying camera calibrtion matrix
    undst = cv2.undistort(image, mtx, dist, None, mtx)
    if write_to_file:
        cv2.imwrite('output_images/undistorted.png', cv2.cvtColor(undst, cv2.COLOR_RGB2BGR))

    # Use a combination of thresholds and color gradients to obtain an image
    # that retains the lane edges while eliminating irrelevant edges as much as possible
    binary_R, s_binary, sxbinary, edges_binary, edges_color = edge_thresholds(undst)
    if write_to_file:
        cv2.imwrite('output_images/edges.png', edges_color)
        R_color = np.dstack((binary_R, binary_R, binary_R))*255
        s_color = np.dstack((s_binary, s_binary, s_binary))*255
        sx_color = np.dstack((sxbinary, sxbinary, sxbinary))*255
        cv2.imwrite('output_images/R.jpg', R_color)
        cv2.imwrite('output_images/s.jpg', s_color)
        cv2.imwrite('output_images/sx.jpg', sx_color)

    # Remove all edges outside the region of interest defined by the vertices
    masked_image, masked_color = region_of_interest(edges_binary, mask_vertices)
    if write_to_file:
        cv2.imwrite('output_images/masked_image.png', masked_color)

    # Transform the masked image to obtain a top view of the image through perspective transform
    topview_image, M, Minv = topview(masked_image)
    topview_color = np.dstack((topview_image, topview_image, topview_image)) * 255
    if write_to_file:
        cv2.imwrite('output_images/topview.png', topview_color)

    # Search the top view image to identify the lane region and fit polynomials to
    # the left, right and middle of the lane
    topview_region, stacked_windows_image, left_fit, right_fit, middle_fit, \
            left_fitx, right_fitx, middle_fitx, ploty = search_around_poly(topview_image)
    if write_to_file:
        cv2.imwrite('output_images/topview_region.png', topview_region)
        cv2.imwrite('output_images/windowed_image.png', cv2.cvtColor(stacked_windows_image, cv2.COLOR_RGB2BGR))

    # Convert polynomial coefficients from image to real coordinate system
    left_fit_real = convert_to_real(left_fit, xm_per_pix, ym_per_pix, 2)
    right_fit_real = convert_to_real(right_fit, xm_per_pix, ym_per_pix, 2)
    middle_fit_real = convert_to_real(middle_fit, xm_per_pix, ym_per_pix, 2)

    # Compute curvature for left, right and middle lines
    left_curverad = measure_curvature(left_fit_real, (image.shape[0])*ym_per_pix)
    right_curverad = measure_curvature(right_fit_real, (image.shape[0])*ym_per_pix)
    middle_curverad = measure_curvature(middle_fit_real, (image.shape[0])*ym_per_pix)

    if False:
        print('Left lane curvature (m) = ', left_curverad)
        print('Right lane curvature (m) = ', right_curverad)
        print('Middle lane curvature (m) = ', middle_curverad)

    # Transform the warped image back into the original image
    warped_back = cv2.warpPerspective(topview_region, Minv, topview_region.shape[1::-1], flags=cv2.INTER_LINEAR)
    if write_to_file:
        cv2.imwrite('output_images/lane_region.png', cv2.cvtColor(warped_back, cv2.COLOR_RGB2BGR))

    # Compute mid point of the lane
    mid_lane = np.array([[middle_fitx[-1], ploty[-1]]], dtype = "float32")
    mid_lane = np.array([mid_lane])
    midpoint_lane = cv2.perspectiveTransform(mid_lane, Minv)

    # Compute midpoint of the image (assume this is the center of the car)
    mid_image = np.array([[image.shape[1]//2, ploty[-1]]], dtype = "float32")
    mid_image = np.array([mid_image])
    midpoint_image = cv2.perspectiveTransform(mid_image, Minv)
    offset = (midpoint_lane[0][0][0] - midpoint_image[0][0][0])*xm_per_pix

    # Combined the lane region markings to the original image
    result = weighted_img(warped_back, image, alpha, 1, gamma)
    if write_to_file:
        cv2.imwrite('output_images/weighted_image.png', cv2.cvtColor(result, cv2.COLOR_RGB2BGR))

    # Write radius of curvature and offset onto the image
    font = cv2.FONT_HERSHEY_SIMPLEX
    radius_loc = (int(0.1*image.shape[0]),int(0.05*image.shape[1]))
    offset_loc = (int(0.1*image.shape[0]),int(0.1*image.shape[1]))
    fontScale = 1.5
    fontColor = (255,255,255)
    lineType = 2
    radius_text = 'Radius of curvature = ' + str(middle_curverad) + '(m)'
    if offset > 0:
        offset_text = 'Vehicle is {0:0.2f}m right of center'.format(offset)
    else:
        offset_text = 'Vehicle is {0:0.2f}m left of center'.format(-offset)

    cv2.putText(result,radius_text, radius_loc, font, fontScale, fontColor, lineType)
    cv2.putText(result, offset_text, offset_loc, font, fontScale, fontColor, lineType)
    if write_to_file:
        cv2.imwrite('output_images/annotated_image.png', cv2.cvtColor(result, cv2.COLOR_RGB2BGR))

    return result
```

### Test Sample Image

```python
#for imgpath in os.listdir("test_images/"):
for imgpath in ["straight_lines1.jpg"]:
    fullpath = os.path.join("test_images", imgpath)
    basename, ext = os.path.splitext(imgpath)
    imgpath_output = basename + '_output' + ext
    fullpath_output = os.path.join("output_images", imgpath_output)

    image = mpimg.imread(fullpath)
    result = process_single_image(image, write_to_file = True)

    cv2.imwrite(fullpath_output, cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
    plt.imshow(result)
```

### Create Video

```python
# Import everything needed to edit/save/watch video clips
from moviepy.editor import VideoFileClip
from IPython.display import HTML
```

```python
# Process video by proecessing the individual images that constitute the video
video_input = 'project_video.mp4'
video_output = 'output_images/project_video_output.mp4'
clip1 = VideoFileClip(video_input).subclip(0, None)
project_clip = clip1.fl_image(process_single_image) #NOTE: this function expects color images!!
%time project_clip.write_videofile(video_output, audio=False)
```

```python
# Embed video
HTML("""

""".format(video_output))
```


# Orbits in space

Now that we know how to describe orbits in the perifocal frame, we need will look at how to move from this frame of reference to a more-general location in space.

This module will cover:
- Geocentric equatorial frame
- Orbital elements
- Gibbs problem
- Gauss problem/Lambert theorem
- Topocentric frame

```python
import numpy as np
%matplotlib inline
from matplotlib import pyplot as plt
```

## Geocentric equatorial frame

In the geocentric equatorial frame, the Earth lies at the center and spins,
where the $X-Y$ plane aligns with the equatorial plane of the planet and
the planet's axis of rotation points in the $Z$ direction of the frame.

Note that this frame does not align with the plane of the ecliptic (the solar plane),
since the Earth's axis is tilted 23.44°.

:::{figure,myclass} fig-gef

Diagram of geocentric equatorial frame.
Original source by [Tfr000](https://en.wikipedia.org/wiki/File:Ra_and_dec_rectangular.png),
shared under [CC BY-SA](https://creativecommons.org/licenses/by-sa/3.0/deed.en)
:::

The $X$ axis of the geocentric equatorial frame always points in the same direction,
towards the vernal equinox. This direction is also referred to as the first point
in Aries, or with ♈, since this pointed to the constellation Aries back
when the Greeks came up with the idea. However, due to the precession of the Earth's tilt,
the direction changes over time, and it now points towards Pisces. In the year 2600,
the vernal equinox will point towards Aquarius.

Due to this precession, for precise orbital determinations we need to specify the
direction of $X$ via the appropriate 50-year **epoch**. We are currently in the
2000 epoch, but in 2025 we will pass into the 2050 epoch.

For an object in space, we can specify its position in the geocentric equatorial
frame via its distance from the center ($r$), the angle away from the $X$ direction
in the equatorial plane is the right ascension ($\alpha$), and the angle above
the plane is the declination ($\delta$).

## Orbital elements

In the perifocal frame, we can specify an object in orbit using the eccentricity $e$,
the specific angular momentum $h$, and the true anomaly $\theta$.
To specify the orientation of this plane in the geocentric equatorial frame,
we need three additional quantities.

:::{figure,myclass} fig-orbital-elements

Diagram showing the orbital elements.
Note that this uses $\nu$ for true anomaly, while we usually use $\theta$ here.
Source: [Lasunncty](https://commons.wikimedia.org/wiki/File:Orbit1.svg) at the English Wikipedia / [CC BY-SA](http://creativecommons.org/licenses/by-sa/3.0/)
:::

The additional three quantities needed are the:
* right ascension $\Omega$, which is the angle between the reference direction $\vec{X}$, or ♈, and the direction of the line of the ascending node $\vec{N}$;
* inclination $i$, the angle of the perifocal frame away from the reference plane; and
* argument of periapse/perigee $\omega$, the angle between the line of ascending node and direction of periapse $\vec{e}$.

The orbit is prograde if $i <$ 90°, and retrograde if $i >$ 90°.

Now, we have the six **orbital elements** needed to uniquely specify an orbit in a two-body system: $h, e, \theta, \Omega, i, \omega$.

### Two-line element sets

Two-line element (TLE) sets, or "elsets", are a standard format used
to encode and distribute the oribital elements of an object in space,
along with other important information. These were originally designed
for a 72-column punchcard format, like what FORTRAN 77 used.

Particular column ranges encode specific information and quantities.
Here is an example for the [NOAA-6 satellite](https://en.wikipedia.org/wiki/NOAA-6).

```
NOAA 6
1 11416U 84123  A 86 50.28438588 0.00000140  00000-0  67960-4 0  5293
2 11416  98.5105  69.3305 0012788  63.2828 296.9658 14.24899292346978
```

Line 0 contains the satellite name (up to 24 characters).
On line 1: by column range:

| Column | Description |
| -------| ----------- |
| 01     | line number (1)       |
| 03-07  | satellite number designated by USSPACECOM (11416) |
| 08     | classification (`U` means unclassified) |
| 10-11  | international designator, launch year (`84` means 1984) |
| 12-14  | international designator, launch number (`123` means 124th launch) |
| 15-17  | international designator, piece (`A` means first object) |
| 19-20  | Epoch year (`86` means 1986) |
| 21-32  | Epoch, Julian day fraction (`50.28438588` converts to 050:06:49:30.94) |
| 34-43  | ballistic coefficient, or 1st time derivative of mean motion (`0.00000140`) |
| 45-52  | 2nd derivative of mean motion, usually blank (`00000-0` = 0.00000) |
| 54-61  | BSTAR, or drag term (`67960-4` means `0.67960e-4`) |
| 63     | ephemeris type (`0`) |
| 65-69  | element number and checksum |

On line 2:

| Column | Description |
| -------| ----------- |
| 01     | line number (2)       |
| 03-07  | satellite number designated by USSPACECOM (11416) |
| 09-16  | inclination in degrees (`98.5105` means 98.5105°) |
| 18-25  | right ascension of the ascending node in degrees (`69.3305` means 69.3305°) |
| 27-33  | eccentricity (`0012788` means 0.0012788) |
| 35-42  | argument of perigee, in degrees (`63.2828` means 63.2828°) |
| 44-51  | mean anomaly, in degrees (`296.9658` means 296.9658°) |
| 53-63  | mean motion, orbits per day (`14.24899292`) |
| 64-68  | revolution number at epoch, in orbits (`34697`) |
| 69     | checksum |

### Calculating orbital elements from state vectors

Given the state vectors $\vec{r}$ and $\vec{v}$, a typical problem is to
calculate the six orbital elements.

First, calculate the magnitudes of relevant vectors:

$$
r = | \vec{r} | \quad v = |\vec{v}| \quad v_r = \frac{\vec{r} \cdot \vec{v}}{r}
$$

Then, calculate each of the orbital elements:

1. Specific angular momentum:

   $$
   \vec{h} = \vec{r} \times \vec{v} \quad h = |\vec{h}|
   $$

2. Inclination:

   $$
   i = \cos^{-1} \left( \frac{h_z}{h} \right)
   $$

3. Right ascension of the ascending node:

   $$
   \begin{align}
   \vec{N} &= \vec{K} \times \vec{h} \\
   N &= \left| \vec{N} \right| \\
   \Omega &= \begin{cases}
     \cos^{-1} \left( \frac{N_x}{N} \right) \quad \text{if } N_Y \geq 0 \;, \\
     360° - \cos^{-1} \left( \frac{N_x}{N} \right) \quad \text{if } N_Y < 0 \;, \\
   \end{cases}
   \end{align}
   $$

4. Eccentricity:

   $$
   \vec{e} = \frac{1}{\mu} \left[ \left( v^2 - \frac{\mu}{r} \right) \vec{r} - r v_r \vec{v} \right] \quad e = | \vec{e} |
   $$

5. Argument of perigee (i.e., argument of periapse):

   $$
   \omega = \begin{cases}
     \cos^{-1} \left( \frac{\vec{N} \cdot \vec{e}}{N e} \right) \quad \text{if } e_Z \geq 0 \;, \\
     360° - \cos^{-1} \left( \frac{\vec{N} \cdot \vec{e}}{N e} \right) \quad \text{if } e_Z < 0 \;, \\
   \end{cases}
   $$

6. True anomaly:

   $$
   \theta = \begin{cases}
   \cos^{-1} \left( \frac{\vec{e} \cdot \vec{r}}{e r} \right) \quad \text{if } v_r \geq 0 \;, \\
     360° - \cos^{-1} \left( \frac{\vec{e} \cdot \vec{r}}{e r} \right) \quad \text{if } v_r < 0 \;, \\
   \end{cases}
   $$

### Calculating state vectors from orbital elements

The opposite from above is calculating the position and velocity state vectors
from the six orbital elements. We can do this by using $h, e, \theta$ to determine
$\vec{r}$ and $\vec{v}$ in the perifocal frame, and then rotating into the
geocentric equatorial frame of reference ($\vec{R}$ and $\vec{V}$).

We can find the overall transformation by combining three rotation matrices:

1. Yaw about the $Z$ axis:

   $$
   R_3 \left( \Omega \right) = \begin{bmatrix}
     \cos \Omega & \sin \Omega & 0 \\
     -\sin \Omega & \cos \Omega & 0 \\
     0 & 0 & 1 \end{bmatrix}
   $$

2. Pitch about the $X_P^{\prime}$ axis:

   $$
   R_1 \left( i \right) = \begin{bmatrix}
     1 & 0 & 0 \\
     0 & \cos i & \sin i \\
     0 & -\sin i & \cos i \end{bmatrix}
   $$

3. Roll about the $Z_P^{\prime}$ axis:

   $$
   R_3 \left( \omega \right) = \begin{bmatrix}
     \cos \omega & \sin \omega & 0 \\
     -\sin \omega & \cos \omega & 0 \\
     0 & 0 & 1 \end{bmatrix}
   $$

4. Calculate the total transformation matrix:

   $$
   \left[ Q_{Xx} \right] = R_3 \left( \omega \right) R_1 (i) R_3 \left( \Omega \right) \;,
   $$

   which involves matrix multiplications.

5. Calculate the new position and velocity vectors in the geocentric equatorial frame:

   $$
   \begin{align}
   \vec{R} &= \left[ Q_{Xx} \right]^{\top} \vec{r} \\
   \vec{V} &= \left[ Q_{Xx} \right]^{\top} \vec{v} \;,
   \end{align}
   $$

   where $\vec{r}$ and $\vec{v}$ are the state vectors in the perifocal frame:

   $$
   \begin{align}
   \vec{r} &= \frac{h^2 / \mu}{1 + e \cos \theta} \begin{bmatrix} \cos \theta & \sin \theta & 0 \end{bmatrix} \\
   \vec{v} &= \frac{\mu}{h} \begin{bmatrix} -\sin \theta & e + \cos \theta & 0 \end{bmatrix}
   \end{align}
   $$

## Gibbs problem

One possible situation is determining the orbit from three sightings / position vectors;
this is the Gibbs problem.

**The problem:** Given $\vec{R}_1, \vec{R}_2, \vec{R}_3$, find the six orbital elements
$h, e, \theta, \omega, \Omega, i$.

:::{figure,myclass} fig-gibbs

Diagram showing Gibbs problem.
:::

Note that the three position vectors must lie in the same plane, otherwise a stable
two-body orbit is not possible.

The goal here is to calculate the velocity vector at location 1, and then with $\vec{R}_1$ and $\vec{V}_1$
the orbital elements can be found using the steps given above:

$$
\vec{V}_1 = \sqrt{ \frac{\mu}{N D} } \left( \frac{\vec{D} \times \vec{R}_1}{r_1} + \vec{S} \right) \;,
$$

where the magnitudes of the position vectors are $r_1 = \left| \vec{R}_1 \right|$, $r_2 = \left| \vec{R}_2 \right|$,
$r_3 = \left| \vec{R}_3 \right|$, and the other vectors are

$$
\begin{align}
\vec{N} &= r_1 \left( \vec{R}_2 \times \vec{R}_3 \right) + r_2 \left( \vec{R}_3 \times \vec{R}_1 \right) + r_3 \left( \vec{R}_1 \times \vec{R}_2 \right) \\
\vec{D} &= \vec{R}_2 \times \vec{R}_3 + \vec{R}_3 \times \vec{R}_1 + \vec{R}_1 \times \vec{R}_2 \\
\vec{S} &= \vec{R}_1 (r_2 - r_3) + \vec{R}_2 (r_3 - r_1) + \vec{R}_3 (r_1 - r_2) \;,
\end{align}
$$

and the magnitudes are $N = \left| \vec{N} \right|$ and $D = \left| \vec{D} \right|$.

:::{warning}
You need to check that $\vec{D} \neq 0$, $\vec{N} \neq 0$, and that $\vec{D} \cdot \vec{N} \geq 0$.
:::

## Lambert's problem

Another fundamental problem is **Lambert's problem**, which is finding the orbit that connects two points in space for a given transfer time.

:::{figure,myclass} fig-lambert

Diagram showing Lambert's problem of orbital transfer between two points.
:::

**Problem:** Given $\vec{R}_1$, $\vec{R}_2$, and $\Delta t$, find the six orbital elements.

Solution steps:

1. Calculate the radial position magnitudes: $r_1 = \left| \vec{R}_1 \right|$ and $r_2 = \left| \vec{R}_2 \right|$.

2. Calculate the change in true anomaly, depending on if the orbit is prograde or retrograde (assume prograde to start):

   $$
   \begin{align}
   \text{if prograde } (i < 90°): \Delta \theta &=
     \begin{cases}
     \cos^{-1} \left( \frac{\vec{R}_1 \cdot \vec{R}_2}{r_1 r_2} \right) \quad \text{if } \left(\vec{R}_1 \times \vec{R}_2 \right)_Z \geq 0 \\
     360° - \cos^{-1} \left( \frac{\vec{R}_1 \cdot \vec{R}_2}{r_1 r_2} \right) \quad \text{if } \left(\vec{R}_1 \times \vec{R}_2 \right)_Z < 0
     \end{cases} \\
   \text{if retrograde } (i > 90°): \Delta \theta &=
     \begin{cases}
     \cos^{-1} \left( \frac{\vec{R}_1 \cdot \vec{R}_2}{r_1 r_2} \right) \quad \text{if } \left(\vec{R}_1 \times \vec{R}_2 \right)_Z < 0 \\
     360° - \cos^{-1} \left( \frac{\vec{R}_1 \cdot \vec{R}_2}{r_1 r_2} \right) \quad \text{if } \left(\vec{R}_1 \times \vec{R}_2 \right)_Z \geq 0
     \end{cases} \\
   \end{align}
   $$

3. Calculate the quantity $A$:

   $$
   A = \sin \Delta \theta \sqrt{ \frac{r_1 r_2}{1 - \cos \Delta \theta}}
   $$

4. Iteratively solve for the universal anomaly parameter $z$:

   $$
   \begin{align}
   y(z) &= r_1 + r_2 + A \frac{z S(z) - 1}{\sqrt{C(z)}} \\
   \sqrt{\mu} \, \Delta t &= \left[ \frac{y(z)}{C(z)}\right]^{3/2} S(z) + A \sqrt{y(z)}
   \end{align}
   $$

   A reasonable initial guess is $z=1$. Then, these equations can be used with an equation solver like `fzero`
   in Matlab or `scipy.optimize.root_scalar` in Python to solve for $z$. If $z < 0$, the orbit is hyperbolic,
   and if $z > 0$ the solution is elliptic.

5. Calculate the universal Lagrange coefficients using the value of $z$:

   $$
   \begin{align}
   f &= 1 - \frac{y(z)}{r_1} \\
   g &= A \sqrt{ \frac{y(z)}{\mu} } \\
   \dot{f} &= \frac{\sqrt{\mu}}{r_1 r_2} \sqrt{ \frac{y(z)}{C(z)}} \left[ z S(z) - 1 \right] \\
   \dot{g} &= 1 - \frac{y(z)}{r_2}
   \end{align}
   $$

   where $C(z)$ and $S(z)$ are the Stumpf functions.

6. Calculate the velocities at each position:

   $$
   \begin{align}
   \vec{V}_1 &= \frac{1}{g} \left( \vec{R}_2 - f \vec{R}_1 \right) \\
   \vec{V}_2 &= \frac{1}{g} \left( \dot{g} \vec{R}_2 - \vec{R}_1 \right) \\
   \end{align}
   $$

7. With either state vector combination (say, $\vec{R}_1$ and $\vec{V}_1$), you can now obtain the orbital elements.

If you find that the orbit is actually retrograde ($i > 90°$), then repeat the process starting at step 2 using the appropriate equation.

## Topocentric frame

The **topocentric frame** describes the location of an orbiting object from a viewing location on the Earth's surface, based on instantaneous geodetic latitude ($\phi$) and longitude ($\theta$).
This will be an instantaneous relationship, since latitude and longitude are fixed but on the rotating surface of the Earth.

**The problem:** find the range and elevation from a viewing point on the surface, given the instantaneous position vector of the orbiting object in the geocentric equatorial frame $\vec{R}_{X}$,
the geodetic latitude $\phi$, and the geodetic longitude $\theta$.

Subscript meanings:

- capital $X$: vector in the geocentric equatorial frame
- lowercase $x$: vector in the topocentric frame

The solution is to find the relative position vector in the geocentric equatorial frame,
then rotate that into the local topocentric frame:

$$
\begin{align}
\vec{\rho}_X &= \vec{R}_X - \vec{R}_{\text{site}} \\
\vec{\rho}_x &= [Q_{Xx}] \vec{\rho}_X \;,
\end{align}
$$

where $\vec{R}_{\text{site}}$ is the position vector from the Earth's center
(i.e., the origin of the geocentric equatorial frame):

$$
\vec{R}_{\text{site}} = \left[ \frac{\bar{R}_e}{\sqrt{1 - (2f - f^2) \sin^2 \phi}} + H \right] \cos \phi
\left\{ \cos \theta \, \vec{I} + \sin \theta \, \vec{J} \right\} + \\
\left[ \frac{\bar{R}_e (1 - f)^2}{\sqrt{1 - (2f - f^2) \sin^2 \phi}} + H \right] \sin \phi \, \vec{K} \;,
$$

where $\bar{R}_e$ = 6378 km is the average radius of the Earth,
$H$ is the altitude/elevation of the observing site, and
$f = 0.003353$ is the oblatedness of the Earth.
Then, $[Q_{Xx}] = R_1 \left( \phi \right) \cdot R_3 \left( \theta \right)$
is the rotation matrix:

$$
\begin{align}
R_1 \left( \phi \right) &= \begin{bmatrix}
1 & 0 & 0 \\
0 & -\sin \phi & \cos \phi \\
0 & \cos \phi  & \sin \phi \end{bmatrix} \\
R_3 \left( \theta \right) &= \begin{bmatrix}
-\sin \theta & \cos \theta & 0 \\
\cos \theta & \sin \theta & 0 \\
0 & 0 & 1 \end{bmatrix} \;.
\end{align}
$$

Then, we can find the range with $\left| \vec{\rho}_x \right|$ and elevation with

$$
a = \sin^{-1} \left( \hat{\rho}_z \right) \quad \hat{\rho} = \frac{\vec{\rho}_x}{\left| \vec{\rho}_x \right|} \;,
$$

where $\hat{\rho}_z$ is the $z$-component of the unit position vector $\hat{\rho}$.

```python

```


# Neuromatch Academy: Week 1, Day 3, Tutorial 1
# Model Fitting: Linear regression with MSE

**Content creators**: Pierre-Étienne Fiquet, Anqi Wu, Alex Hyafil with help from Byron Galbraith

**Content reviewers**: Lina Teichmann, Saeed Salehi, Patrick Mineault,  Ella Batty, Michael Waskom

___
#Tutorial Objectives

This is Tutorial 1 of a series on fitting models to data. We start with simple linear regression, using least squares optimization (Tutorial 1) and Maximum Likelihood Estimation (Tutorial 2). We will use bootstrapping to build confidence intervals around the inferred linear model parameters (Tutorial 3). We'll finish our exploration of regression models by generalizing to multiple linear regression and polynomial regression (Tutorial 4). We end by learning how to choose between these various models. We discuss the bias-variance trade-off (Tutorial 5) and Cross Validation for model selection (Tutorial 6).

In this tutorial, we will learn how to fit simple linear models to data.
- Learn how to calculate the mean-squared error (MSE)
- Explore how model parameters (slope) influence the MSE
- Learn how to find the optimal model parameter using least-squares optimization

---

**acknowledgements:**
- we thank Eero Simoncelli, much of today's tutorials are inspired by exercises asigned in his mathtools class.

---
# Setup

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

    Matplotlib is building the font cache; this may take a moment.

```python
#@title Figure Settings
import ipywidgets as widgets       # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions

def plot_observed_vs_predicted(x, y, y_hat, theta_hat):
    """ Plot observed vs predicted data

    Args:
        x (ndarray): observed x values
    y (ndarray): observed y values
    y_hat (ndarray): predicted y values

    """
    fig, ax = plt.subplots()
    ax.scatter(x, y, label='Observed')  # our data scatter plot
    ax.plot(x, y_hat, color='r', label='Fit')  # our estimated model
    # plot residuals
    ymin = np.minimum(y, y_hat)
    ymax = np.maximum(y, y_hat)
    ax.vlines(x, ymin, ymax, 'g', alpha=0.5, label='Residuals')
    ax.set(
        title=
        fr"$\hat{{\theta}}$ = {theta_hat:0.2f}, MSE = {mse(x, y, theta_hat):.2f}",
        xlabel='x',
        ylabel='y')
    ax.legend()
```

---
# Section 1: Mean Squared Error (MSE)

```python
#@title Video 1: Linear Regression & Mean Squared Error
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="HumajfjJ37E", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

**Linear least squares regression** is an old but gold  optimization procedure that we are going to use for data fitting. Least squares (LS) optimization problems are those in which the objective function is a quadratic function of the
parameter(s) being optimized.

Suppose you have a set of measurements, $y_{n}$ (the "dependent" variable) obtained for different input values, $x_{n}$ (the "independent" or "explanatory" variable). Suppose we believe the measurements are proportional to the input values, but are corrupted by some (random) measurement errors, $\epsilon_{n}$, that is:

$$y_{n}= \theta x_{n}+\epsilon_{n}$$

for some unknown slope parameter $\theta.$ The least squares regression problem uses **mean squared error (MSE)** as its objective function, it aims to find the value of the parameter $\theta$ by minimizing the average of squared errors:

\begin{align}
\min _{\theta} \frac{1}{N}\sum_{n=1}^{N}\left(y_{n}-\theta x_{n}\right)^{2}
\end{align}

We will now explore how MSE is used in fitting a linear regression model to data. For illustrative purposes, we will create a simple synthetic dataset where we know the true underlying model. This will allow us to see how our estimation efforts compare in uncovering the real model (though in practice we rarely have this luxury).

First we will generate some noisy samples $x$ from [0, 10) along the line $y = 1.2x$ as our dataset we wish to fit a model to.

```python
# @title

# @markdown Execute this cell to generate some simulated data

# setting a fixed seed to our random number generator ensures we will always
# get the same psuedorandom number sequence
np.random.seed(121)

# Let's set some parameters
theta = 1.2
n_samples = 30

# Draw x and then calculate y
x = 10 * np.random.rand(n_samples)  # sample from a uniform distribution over [0,10)
noise = np.random.randn(n_samples)  # sample from a standard normal distribution
y = theta * x + noise

# Plot the results
fig, ax = plt.subplots()
ax.scatter(x, y)  # produces a scatter plot
ax.set(xlabel='x', ylabel='y');
```

Now that we have our suitably noisy dataset, we can start trying to estimate the underlying model that produced it. We use MSE to evaluate how successful a particular slope estimate $\hat{\theta}$ is for explaining the data, with the closer to 0 the MSE is, the better our estimate fits the data.

## Exercise 1: Compute MSE

In this exercise you will implement a method to compute the mean squared error for a set of inputs $x$, measurements $y$, and slope estimate $\hat{\theta}$. We will then compute and print the mean squared error for 3 different choices of theta

```python
def mse(x, y, theta_hat):
    """Compute the mean squared error

    Args:
      x (ndarray): An array of shape (samples,) that contains the input values.
      y (ndarray): An array of shape (samples,) that contains the corresponding
        measurement values to the inputs.
      theta_hat (float): An estimate of the slope parameter

    Returns:
      float: The mean squared error of the data with the estimated parameter.
    """
    # Compute the estimated y
    y_hat = theta_hat * x

    # Compute mean squared error
    mse = np.mean((y - y_hat)**2)

    return mse

# Uncomment below to test your function
theta_hats = [0.75, 1.0, 1.5]
for theta_hat in theta_hats:
    print(f"theta_hat of {theta_hat} has an MSE of {mse(x, y, theta_hat):.2f}")
```

    theta_hat of 0.75 has an MSE of 9.08
    theta_hat of 1.0 has an MSE of 3.01
    theta_hat of 1.5 has an MSE of 4.52

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial1_Solution_12a57de0.py)

The result should be:

theta_hat of 0.75 has an MSE of 9.08\
theta_hat of 1.0 has an MSE of 3.0\
theta_hat of 1.5 has an MSE of 4.52

We see that $\hat{\theta} = 1.0$ is our best estimate from the three we tried. Looking just at the raw numbers, however, isn't always satisfying, so let's visualize what our estimated model looks like over the data.

```python
# @title

# @markdown Execute this cell to visualize estimated models

fig, axes = plt.subplots(ncols=3, figsize=(18, 4))
for theta_hat, ax in zip(theta_hats, axes):

    # True data
    ax.scatter(x, y, label='Observed')  # our data scatter plot

    # Compute and plot predictions
    y_hat = theta_hat * x
    ax.plot(x, y_hat, color='r', label='Fit')  # our estimated model

    ax.set(title=
           fr'$\hat{{\theta}}$= {theta_hat}, MSE = {mse(x, y, theta_hat):.2f}',
           xlabel='x',
           ylabel='y')

axes[0].legend();
```

## Interactive Demo: MSE Explorer

Using an interactive widget, we can easily see how changing our slope estimate changes our model fit. We display the **residuals**, the differences between observed and predicted data, as line segments between the data point (observed response) and the corresponding predicted response on the model fit line.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

@widgets.interact(theta_hat=widgets.FloatSlider(1.0, min=0.0, max=2.0))
def plot_data_estimate(theta_hat):
    y_hat = theta_hat * x
    plot_observed_vs_predicted(x, y, y_hat, theta_hat)
```

    interactive(children=(FloatSlider(value=1.0, description='theta_hat', max=2.0), Output()), _dom_classes=('widg…

While visually exploring several estimates can be instructive, it's not the most efficient for finding the best estimate to fit our data. Another technique we can use is choose a reasonable range of parameter values and compute the MSE at several values in that interval. This allows us to plot the error against the parameter value (this is also called an **error landscape**, especially when we deal with more than one parameter). We can select the final $\hat{\theta}$ ($\hat{\theta}_{MSE}$) as the one which results in the lowest error.

```python
# @title

# @markdown Execute this cell to loop over theta_hats, compute MSE, and plot results

# Loop over different thetas, compute MSE for each
theta_hat_grid = np.linspace(-2.0, 4.0)
errors = np.zeros(len(theta_hat_grid))
for i, theta_hat in enumerate(theta_hat_grid):
    errors[i] = mse(x, y, theta_hat)

# Find theta that results in lowest error
best_error = np.min(errors)
theta_hat = theta_hat_grid[np.argmin(errors)]

# Plot results
fig, ax = plt.subplots()
ax.plot(theta_hat_grid, errors, '-o', label='MSE', c='C1')
ax.axvline(theta, color='g', ls='--', label=r"$\theta_{True}$")
ax.axvline(theta_hat, color='r', ls='-', label=r"$\hat{{\theta}}_{MSE}$")
ax.set(
    title=
    fr"Best fit: $\hat{{\theta}}$ = {theta_hat:.2f}, MSE = {best_error:.2f}",
    xlabel=r"$\hat{{\theta}}$",
    ylabel='MSE')
ax.legend()
```

We can see that our best fit is $\hat{\theta}=1.18$ with an MSE of 1.45. This is quite close to the original true value $\theta=1.2$!

---
# Section 2: Least-squares optimization

While the approach detailed above (computing MSE at various values of $\hat\theta$) quickly got us to a good estimate, it still relied on evaluating the MSE value across a grid of hand-specified values. If we didn't pick a good range to begin with, or with enough granularity, we might miss the best possible estimator. Let's go one step further, and instead of finding the minimum MSE from a set of candidate estimates, let's solve for it analytically.

We can do this by minimizing the cost function. Mean squared error is a convex objective function, therefore we can compute its minimum using calculus. Please see video or appendix for this derivation! After computing the minimum, we find that:

\begin{align}
\hat\theta = \frac{\vec{x}^\top \vec{y}}{\vec{x}^\top \vec{x}}
\end{align}

This is known as solving the normal equations. For different ways of obtaining the solution, see the notes on [Least Squares Optimization](https://www.cns.nyu.edu/~eero/NOTES/leastSquares.pdf) by Eero Simoncelli.

### Exercise 2: Solve for the Optimal Estimator

In this exercise, you will write a function that finds the optimal $\hat{\theta}$ value using the least squares optimization approach (the equation above) to solve MSE minimization. It shoud take arguments $x$ and $y$ and return the solution $\hat{\theta}$.

We will then use your function to compute $\hat{\theta}$ and plot the resulting prediction on top of the data.

```python
def solve_normal_eqn(x, y):
    """Solve the normal equations to produce the value of theta_hat that minimizes
      MSE.

      Args:
      x (ndarray): An array of shape (samples,) that contains the input values.
      y (ndarray): An array of shape (samples,) that contains the corresponding
        measurement values to the inputs.

    Returns:
      float: the value for theta_hat arrived from minimizing MSE
    """
    # Compute theta_hat analytically
    theta_hat = np.sum(x * y)/ np.sum(x**2)
    # or
    # theta_hat = (x.T @ y) / (x.T @ x) # Also works, note as this are 1d .T doesn't have an effect

    return theta_hat

# Uncomment below to test your function
theta_hat = solve_normal_eqn(x, y)
y_hat = theta_hat * x
plot_observed_vs_predicted(x, y, y_hat, theta_hat)
```

The method used above is an alternative for the same solution, suggested in the video:

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial1_Solution_7a89ba24.py)

*Example output:*

We see that the analytic solution produces an even better result than our grid search from before, producing $\hat{\theta} = 1.21$ with MSE = 1.43!

---
# Summary

- Linear least squares regression is an optimization procedure that can be used for data fitting:
    - Task: predict a value for $y$ given $x$
    - Performance measure: $\textrm{MSE}$
    - Procedure: minimize $\textrm{MSE}$ by solving the normal equations
- **Key point**: We fit the model by defining an *objective function* and minimizing it.
- **Note**: In this case, there is an *analytical* solution to the minimization problem and in practice, this solution can be computed using *linear algebra*. This is *extremely* powerful and forms the basis for much of numerical computation throughout the sciences.

---
# Appendix

## Least Squares Optimization Derivation

We will outline here the derivation of the least squares solution.

We first set the derivative of the error expression with respect to $\theta$ equal to zero,

\begin{align}
\frac{d}{d\theta}\frac{1}{N}\sum_{i=1}^N(y_i - \theta x_i)^2 = 0 \\
\frac{1}{N}\sum_{i=1}^N-2x_i(y_i - \theta x_i) = 0
\end{align}

where we used the chain rule. Now solving for $\theta$, we obtain an optimal value of:

\begin{align}
\hat\theta = \frac{\sum_{i=1}^N x_i y_i}{\sum_{i=1}^N x_i^2}
\end{align}

Which we can write in vector notation as:

\begin{align}
\hat\theta = \frac{\vec{x}^\top \vec{y}}{\vec{x}^\top \vec{x}}
\end{align}

This is known as solving the *normal equations*. For different ways of obtaining the solution, see the notes on [Least Squares Optimization](https://www.cns.nyu.edu/~eero/NOTES/leastSquares.pdf) by Eero Simoncelli.