## Long Run Mean Growth Rates for Discount and Return Processes

#### Qingyin Ma, John Stachurski and Alexis Akira Toda

```python
import numpy as np
from scipy.linalg import eig, eigvals
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib import cm
import quantecon as qe
```

```python
def tauchen(mu, rho, sigma_u, N=5, n_std=3):
    """
    Using Tauchen's method to discretization to compute the Markov
    matrix associated with the linear Gaussian AR(1) process

        y_{t+1} = mu + rho * y_t + u_{t+1},
        {u_t} ~iid N(0, sigma_u^2).

    Parameters
    ----------
    mu  : scalar(float)
        The constant term of {y_t}
    rho : scalar(float)
        The autocorrelation coefficient
    sigma_u : scalar(float)
        The standard deviation of {u_t}
    N : scalar(int), optional(default=7)
        The number of states to use in the approximation
    n_std : scalar(int), optional(default=3)
        The number of standard deviations to approximate out to

    Returns
    -------
    states : array_like(float, ndim=1)
        The state space of the discretized process
    P : array_like(float, ndim=2)
        The Markov transition matrix where P[i, j] is the probability
        of transitioning from x[i] to x[j]

    """
    # CDF of the normal distribution
    F = norm(loc=0, scale=sigma_u).cdf

    # the stationary standard deviation of {y_t}
    std_y = np.sqrt(sigma_u**2 / (1-rho**2))

    # distance from the center
    x_bar = n_std * std_y

    # discretized state space (mean normalized)
    x = np.linspace(-x_bar, x_bar, N)

    step = x[1] - x[0]
    half_step = 0.5 * step
    P = np.empty((N, N))

    for i in range(N):
        P[i, 0] = F(x[0] - rho * x[i] + half_step)
        P[i, N-1] = 1 - F(x[N-1] - rho * x[i] - half_step)
        for j in range(1, N-1):
            z = x[j] - rho * x[i]
            P[i, j] = F(z + half_step) - F(z - half_step)

    # renormalize the matrix to avoid rounding errors
    row_sums = P.sum(axis=1)
    P = P / row_sums[:, np.newaxis]

    # center process around its mean: mu / (1 - rho)
    states = mu / (1 - rho) + x

    return states, P
```

## Random Discounting

```python
class RandomDiscounting:
    """
    The AR(1) random discounting model.
    """
    def __init__(self,
                 μ=0.99,    # stationary mean of {β_t}
                 ρ=0.992,   # autocorrelation coefficient of {β_t}
                 σ=0.006,   # standard deviation of {β_t}
                 N=15):     # number of states of the discretized {β_t} process
        self.μ, self.ρ, self.σ, self.N = μ, ρ, σ, N

        # discretize {β_t} process
        #self.Z_β, self.P = tauchen((1 - ρ) * μ, ρ, np.sqrt(1 - ρ**2) * σ, N=N)
        mc = qe.rouwenhorst(N, (1 - ρ) * μ, np.sqrt(1 - ρ**2) * σ, ρ)
        self.Z_β, self.P = mc.state_values, mc.P

    def Gβ(self):
        """
        Evaluate the long-run average growth rate of {β_t}
        """
        Z_β, P = self.Z_β, self.P

        D_β = np.diag(Z_β)
        L_β = P @ D_β
        G_β = max(np.abs(eigvals(L_β)))  # G_β = r(L_β)
        return G_β
```

```python
rd = RandomDiscounting()
```

```python
rd.Gβ()
```

    0.9943155894974464

```python
rd.Z_β
```

    array([0.96755006, 0.97075719, 0.97396433, 0.97717146, 0.9803786 ,
           0.98358573, 0.98679287, 0.99      , 0.99320713, 0.99641427,
           0.9996214 , 1.00282854, 1.00603567, 1.00924281, 1.01244994])

```python
M = 50
mu_seq = np.linspace(0.94, 1.04, M)
out = np.empty(M)
for i in range(M):
    rd = RandomDiscounting(μ=mu_seq[i])
    out[i] = rd.Gβ()

fig, ax = plt.subplots()
ax.plot(mu_seq, out)
plt.show()
```

A function for generating contour plots.

```python
def contour_Gβ(ρ_min=0.96,       # grid point of minimum ρ value
               ρ_max=0.99,    # grid point of maximum ρ value
               ρ_size=50,       # size of grid points for ρ
               σ_min=0.01,    # grid point of minimum σ value
               σ_max=0.015,       # grid point of maximum σ value
               σ_size=50):      # number of grid points for σ
    """
    Evaluate G_β for a class of ρ and σ values.
    """
    ρ_grid = np.linspace(ρ_min, ρ_max, ρ_size)  # grid points for ρ
    σ_grid = np.linspace(σ_min, σ_max, σ_size)  # grid points for σ

    Gβ_matrix = np.empty((ρ_size, σ_size))  # create empty matrix to store G_β values

    for (i_ρ, ρ) in enumerate(ρ_grid):
        for (i_σ, σ) in enumerate(σ_grid):
            rd = RandomDiscounting(ρ=ρ, σ=σ)
            Gβ_matrix[i_ρ, i_σ] = rd.Gβ()

    return Gβ_matrix, ρ_grid, σ_grid
```

```python
Gβ_matrix, ρ_grid, σ_grid = contour_Gβ(ρ_size=12, σ_size=12)
```

#### Make contour plot of G_β

```python
fig, ax = plt.subplots(figsize=(10,6))

Xβ, Yβ = np.meshgrid(ρ_grid, σ_grid)

# hide high values in Gβ_matrix
Zβ = Gβ_matrix
Zβ[Zβ > 1.5] = np.nan

CSβ_f=ax.contourf(Xβ, Yβ, Zβ.T, alpha=0.4)
plt.colorbar(CSβ_f, ax=ax)

levels=[1.0]

CSβ = ax.contour(Xβ, Yβ, Gβ_matrix.T, colors='k',
                 levels=levels)

ax.clabel(CSβ, inline=1, fontsize=15, fmt="%1.2f")

ax.set_xlabel("$\\rho$", fontsize=20)
ax.set_ylabel("$\sigma$", fontsize=20)

plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

plt.savefig('gb.pdf')
plt.show()
```

```python

```

# Verifying the Stability Conditions

### The model is as follows:

> $\log R_t = \mu_t + \sigma_t \zeta_t$,
 $\, \{\zeta_t\} \overset{iid}{\sim} N(0,1)$

 $\mu_t = \bar{\mu} (1 - \rho_\mu) + \rho_\mu \mu_{t-1}
    + \upsilon_t^\mu$,
  $\, (\upsilon_t^\mu) \overset{iid}{\sim} N(0, \delta_\mu^2)$

 $\log \sigma_t = \bar{\sigma} (1 - \rho_\sigma)
    + \rho_\sigma \log \sigma_{t-1}
    + \upsilon_t^\sigma$,
  $\, \{\upsilon_t^\sigma\} \overset{iid}{\sim} N(0, \delta_\sigma^2)$

We use Tauchen's method to discretize $\{\mu_t\}$ and $\{\sigma_t\}$ into finite-state Markov chains,

with state spaces $\mathsf{Z_\mu}$ and $\mathsf{Z_\sigma}$, respectively.

### Define the following:

- $\mathsf{Z}_\mu := \{i_1, \cdots, i_M \},
   \quad
   \mathsf{Z}_\sigma := \{ j_1, \cdots, j_N \}$

- $D_\mu := \text{diag}
            \left(
                \text{e}^{i_1}, \cdots, \text{e}^{i_M}
            \right),
   \quad
   D_\sigma := \text{diag}
               \left(
                   \text{e}^{\, j_1^2 / 2},
                   \cdots,
                   \text{e}^{\, j_N^2 / 2}
               \right)$

- $V_\mu := \left(
                \text{e}^{(1 - \gamma) i_1},
                \cdots,
                \text{e}^{(1 - \gamma) i_M}
            \right)',
   \quad
   V_\sigma := \left(
                   \text{e}^{(1 - \gamma)^2 j_1^2 / 2},
                   \cdots,
                   \text{e}^{(1 - \gamma)^2 j_N^2 / 2}
               \right)'$

### Sufficient conditions for stability:

\begin{equation}
    \beta \, r(P_\mu D_\mu) \cdot r(P_\sigma D_\sigma) < 1
\end{equation}

\begin{equation}
    \kappa := \beta \, \| P_\mu V_\mu \| \cdot \| P_\sigma V_\sigma \| < 1
\end{equation}

\begin{equation}
    r(P_\mu D_\mu) \cdot r(P_\sigma D_\sigma) \cdot \kappa^{1 / \gamma} < 1
\end{equation}

where
- $\| \mathbf{x} \|$ denotes the largest modulus of all the elements in vector $\mathbf{x}$.
- $r(A)$ is the spectral radius of matrix $A$, i.e., the largest modulus of all the eigenvalues of A.

```python

```

## Model I. Constant $\{ \mu_t\}$ process

The model reduces to:

> $\log R_t = \mu + \sigma_t \zeta_t$,
 $\, \{\zeta_t\} \overset{iid}{\sim} N(0,1)$

 $\log \sigma_t = \bar{\sigma} (1 - \rho_\sigma)
    + \rho_\sigma \log \sigma_{t-1}
    + \upsilon_t^\sigma$,
  $\, \{\upsilon_t^\sigma\} \overset{iid}{\sim} N(0, \delta_\sigma^2)$

- $\mu = \bar \mu$ is the stationary mean of the AR(1) process $\{\mu_t\}$.

```python

```

```python
class IFP_StochasticVolatility:
    """
    Model I. Stochastic wealth returns with stochastic volatility.
    """
    def __init__(self, β=0.95,    # the discount factor
                 γ=1.5,           # the relative risk aversion coefficient
                 μ=0.0281,        # mean parameter related to {R_t}
                 σ_bar=-3.2556,   # stationary mean of the AR(1) process {σ_t}
                 ρ_σ=0.2895,      # autoregressive coefficient of the AR(1) process {σ_t}
                 δ_σ=0.1896,      # standard deviation of the {υσ_t} process
                 K=5,             # number of states of the discretized {σ_t} process
                 verbose=False):
        self.β, self.γ, self.K, self.verbose = β, γ, K, verbose
        self.μ, self.σ_bar, self.ρ_σ, self.δ_σ = μ, σ_bar, ρ_σ, δ_σ

        # discretize {σ_t} process
        # self.Z_σ, self.P_σ = tauchen(self.σ_bar*(1-self.ρ_σ), self.ρ_σ, self.δ_σ, N=self.K)
        mc = qe.rouwenhorst(K, σ_bar*(1-ρ_σ), δ_σ, ρ_σ)
        self.Z_σ, self.P_σ = mc.state_values, mc.P

        self.Z_σ = np.exp(self.Z_σ)

    def stability_test(self):
        """
        Stability test of the model economy.
        """
        β, γ, verbose = self.β, self.γ, self.verbose
        μ, Z_σ, P_σ = self.μ, self.Z_σ, self.P_σ

        # ===== long-run average growth rate of {R_t} ===== #
        D_R = np.diag(np.exp(μ + Z_σ**2 / 2))
        if D_R.max() == np.inf:
            G_R = np.nan
        else:
            L_R = np.matmul(P_σ, D_R)        # L_R = P D_R
            G_R = max(np.abs(eig(L_R)[0]))   # G_R = r(L_R)

        # ============= stability test ============= #
        D_μ = np.exp(μ)
        D_σ = np.diag(np.exp(Z_σ**2 / 2.))

        V_μ = np.exp((1. - γ) * μ)
        V_σ = np.exp(((1. - γ)**2) * (Z_σ**2) / 2.)

        if D_σ.max()==np.inf or V_σ.max()==np.inf:
            lhs_cond1, lhs_cond2, lhs_cond3 = np.nan, np.nan, np.nan
        else:
            # r(P_μ D_μ) r(P_σ D_σ)
            spec_rad = D_μ * max(eig(np.matmul(P_σ, D_σ))[0])
            # β r(P_μ D_μ) r(P_σ D_σ)
            lhs_cond1 = β * spec_rad
            # κ := β \|P_μ V_μ \| \|P_σ V_σ\|
            lhs_cond2 = β * V_μ * max(np.matmul(P_σ, V_σ))
            # r(P_μ D_μ) r(P_σ D_σ) κ^{1/γ}
            lhs_cond3 = spec_rad * (lhs_cond2**(1/γ))

        # sufficient conditions for stability
        cond1 = (lhs_cond1 < 1.)
        cond2 = (lhs_cond2 < 1.)
        cond3 = (lhs_cond3 < 1.)
        stability = (cond1 and cond2 and cond3)

        if verbose:
            if cond1 and cond2 and cond3:
                print("")
                print("Stability conditions are verified!")
                print("")
                print("Ready for the next step!")
                print("")
            else:
                print("")
                print("Error!!!")
                print("")
                print("Stability conditions are violated!")
                print("")
                print("Adjust parameter values!")
                print("")
            print("------------------------------------")
            print("G_R = ", G_R)
            print("")
            print("Condition 1: ", lhs_cond1)
            print("")
            print("Condition 2: ", lhs_cond2)
            print("")
            print("Condition 3: ", lhs_cond3)
            print("------------------------------------")
            print("")

        return G_R, stability
```

```python

```

```python
# stability range
def stability_range_sv(ρ_σ_min=0.0,     # grid point of minimum ρ_σ value
                        ρ_σ_max=0.99,   # grid point of maximum ρ_σ value
                        ρ_σ_size=50,    # size of grid points for ρ_σ
                        δ_σ_min=1e-3,   # grid point of minimum δ_σ value
                        δ_σ_max=1.4,    # grid point of minimum δ_σ value
                        δ_σ_size=50):   # size of grid points for δ_σ

    ρ_σ_grid = np.linspace(ρ_σ_min, ρ_σ_max, ρ_σ_size)  # grid points for ρ_σ
    δ_σ_grid = np.linspace(δ_σ_min, δ_σ_max, δ_σ_size)  # grid points for δ_σ

    info_matrix = np.empty((ρ_σ_size, δ_σ_size))  # create empty matrix to store stability test results
    GR_matrix = np.empty((ρ_σ_size, δ_σ_size))    # create empty matrix to store G_R values

    for (i_ρ, ρ) in enumerate(ρ_σ_grid):
        for (i_σ, σ) in enumerate(δ_σ_grid):
            ifp_svr = IFP_StochasticVolatility(ρ_σ=ρ, δ_σ=σ)
            GR_matrix[i_ρ,i_σ], info_matrix[i_ρ,i_σ] = ifp_svr.stability_test()

    return GR_matrix, info_matrix, ρ_σ_grid, δ_σ_grid

GR_matrix, info_matrix, ρ_σ_grid, δ_σ_grid = stability_range_sv(ρ_σ_size=100, δ_σ_size=100)
```

    C:\Users\Public\software\Anaconda3\lib\site-packages\ipykernel_launcher.py:55: RuntimeWarning: overflow encountered in cdouble_scalars
    C:\Users\Public\software\Anaconda3\lib\site-packages\ipykernel_launcher.py:31: RuntimeWarning: overflow encountered in exp
    C:\Users\Public\software\Anaconda3\lib\site-packages\ipykernel_launcher.py:40: RuntimeWarning: overflow encountered in exp
    C:\Users\Public\software\Anaconda3\lib\site-packages\ipykernel_launcher.py:43: RuntimeWarning: overflow encountered in exp

```python

```

```python
# find the threshold of stability range
idx = np.minimum(np.sum(info_matrix, axis=1), len(δ_σ_grid)-1)
idx = idx.astype(int)     # threshold index of the δ_σ grid points
δ_thresh = δ_σ_grid[idx]  # threshold δ_σ values

# Plot the stability range
fig, ax = plt.subplots(figsize=(8,6))

ax.plot(ρ_σ_grid, δ_thresh, label='stability threshold', lw=2.5)

ax.set_xlabel(r'$\rho_\sigma$', size=20)
ax.set_ylabel(r'$\delta_\sigma$', size=20)
ax.set_xlim((0.0, 1.0))
ax.set_ylim((0.0, 1.4))
#ax.set_title('Model 1: γ=%i' %γ, size=20)
plt.xticks(fontsize=15)
plt.yticks([0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4], fontsize=15)

font = {'family': 'sans-serif',
        'color':  'black',
        'weight': 'normal',
        'size': 20,
        }

ax.text(0.35, 0.5, 'stable', fontdict=font)
#ax.text(0.55, 1.0, 'unstable', fontdict=font)

ax.plot([0.2895], [0.1896], marker='o', markersize=5, color="blue")
ax.text(0.3, 0.1, "$(0.2895, \, 0.1896)$", family="sans-serif", size=15)

ax.fill_between(ρ_σ_grid, δ_thresh, color='lightyellow')

ax.legend(fontsize=15)

plt.savefig('sv_stb.pdf')

plt.show()
```

```python

```

```python
# ------------------------------------- #
#     make contour plot of G_R
# ------------------------------------- #
fig, ax = plt.subplots(figsize=(10,6))

X1, Y1 = np.meshgrid(ρ_σ_grid, δ_σ_grid)

# hide high values in GR_matrix
Z1 = GR_matrix
Z1[Z1 > 1.2] = np.nan

β = IFP_StochasticVolatility().β

CS1_f = ax.contourf(X1, Y1, β*Z1.T, alpha=0.4)
plt.colorbar(CS1_f, ax=ax)

levels=[1.00]

CS1 = ax.contour(X1, Y1, β*Z1.T, colors='k', levels=levels)
ax.clabel(CS1, inline=1, fontsize=15, fmt="%1.2f")

ax.plot([0.2895], [0.1896], marker='o', markersize=5, color="black")
ax.text(0.3, 0.1, "$(0.2895, \, 0.1896)$", family="sans-serif", size=15)

ax.set_xlabel("$\\rho_\sigma$", fontsize=20)
ax.set_ylabel("$\delta_\sigma$", fontsize=20)

plt.xticks([0.0,0.2,0.4,0.6,0.8,1.0], fontsize=15)
plt.yticks([0.0,0.2,0.4,0.6,0.8,1.0,1.2,1.4], fontsize=15)

plt.savefig('sv_gr.pdf')

plt.show()
```

```python

```

## Model II: Constant $\{ \sigma_t \}$ Process

The model reduces to:

> $\log R_t = \mu_t + \sigma \zeta_t$,
$(\zeta_t) \overset{iid}{\sim} N(0,1)$

> $\mu_t = \bar{\mu} (1 - \rho_\mu) + \rho_\mu \mu_{t-1}
    + \upsilon_t^\mu$,
  $\, (\upsilon_t^\mu) \overset{iid}{\sim} N(0, \delta_\mu^2)$

- $\sigma = e^{\bar \sigma + \delta_\sigma^2 / 2(1 - \rho_\sigma^2)}$:
the stationary mean of the AR(1) process $\{\sigma_t\}$

```python

```

```python
class IFP_MeanPersistence:
    """
    Model II. Stochastic wealth returns with mean persistence.
    """
    def __init__(self, β=0.95,  # the discount factor
                 γ=1.5,         # the relative risk aversion coefficient
                 σ=0.0393,      # standard deviation parameter related to {R_t}
                 μ_bar=0.0281,  # stationary mean of the AR(1) process {μ_t}
                 ρ_μ=0.5722,    # autoregressive coefficient of the AR(1) process {μ_t}
                 δ_μ=0.0067,    # standard deviation of the {υμ_t} process
                 M=5,           # number of states of the discretized {μ_t} process
                 verbose=False):
        self.β, self.γ, self.M, self.verbose = β, γ, M, verbose
        self.σ, self.μ_bar, self.ρ_μ, self.δ_μ = σ, μ_bar, ρ_μ, δ_μ

        # discretize {μ_t} process
        #self.Z_μ, self.P_μ = tauchen(self.μ_bar*(1-self.ρ_μ), self.ρ_μ, self.δ_μ, N=self.M)
        mc = qe.rouwenhorst(M, μ_bar*(1-ρ_μ), δ_μ, ρ_μ)
        self.Z_μ, self.P_μ = mc.state_values, mc.P

    def stability_test(self):
        """
        Stability test of the model economy.
        """
        β, γ, M, verbose = self.β, self.γ, self.M, self.verbose
        σ, Z_μ, P_μ = self.σ, self.Z_μ, self.P_μ

        # ===== long-run average growth rate of {R_t} ===== #
        D_R = np.diag(np.exp(Z_μ + σ**2 / 2))

        if D_R.max() == np.inf:
            G_R = np.nan
        else:
            L_R = np.matmul(P_μ, D_R)        # L_R = P D_R
            G_R = max(np.abs(eig(L_R)[0]))   # G_R = r(L_R)

        # ============= stability test ============= #
        D_σ = np.exp(σ**2 / 2.)
        D_μ = np.diag(np.exp(Z_μ))

        V_σ = np.exp(((1. - γ)**2) * (σ**2) / 2.)
        V_μ = np.exp((1. - γ) * Z_μ)

        if D_μ.max()==np.inf or V_μ.max()==np.inf:
            lhs_cond1, lhs_cond2, lhs_cond3 = np.nan, np.nan, np.nan
        else:
            # r(P_R D_R) = r(P_μ D_μ) r(P_σ D_σ)
            spec_rad = D_σ * max(eig(np.matmul(P_μ, D_μ))[0])
            # β r(P_μ D_μ) r(P_σ D_σ)
            lhs_cond1 = β * spec_rad
            # κ := β \|P_μ V_μ \| \|P_σ V_σ\|
            lhs_cond2 = β * V_σ * max(np.matmul(P_μ, V_μ))
            # r(P_μ D_μ) r(P_σ D_σ) κ^{1/γ}
            lhs_cond3 = spec_rad * (lhs_cond2**(1/γ))

        # sufficient conditions for stability
        cond1 = (lhs_cond1 < 1.)
        cond2 = (lhs_cond2 < 1.)
        cond3 = (lhs_cond3 < 1.)
        stability = (cond1 and cond2 and cond3)

        if verbose:
            if cond1 and cond2 and cond3:
                print("")
                print("Stability conditions are verified!")
                print("")
                print("Ready for the next step!")
                print("")
            else:
                print("")
                print("Error!!!")
                print("")
                print("Stability conditions are violated!")
                print("")
                print("Adjust parameter values!")
                print("")
            print("------------------------------------")
            print("G_R = ", G_R)
            print("")
            print("Condition 1: ", lhs_cond1)
            print("")
            print("Condition 2: ", lhs_cond2)
            print("")
            print("Condition 3: ", lhs_cond3)
            print("------------------------------------")
            print("")

        return G_R, stability
```

```python

```

```python
# stability range
def stability_range_mp(ρ_μ_min=0.0,    # grid point of minimum ρ_μ value
                       ρ_μ_max=0.99,   # grid point of maximum ρ_μ value
                       ρ_μ_size=50,    # size of grid points for ρ_μ
                       δ_μ_min=1e-3,   # grid point of minimum δ_μ value
                       δ_μ_max=0.14,   # grid point of minimum δ_μ value
                       δ_μ_size=50):   # size of grid points for δ_μ

    ρ_μ_grid = np.linspace(ρ_μ_min, ρ_μ_max, ρ_μ_size)  # grid points for ρ_μ
    δ_μ_grid = np.linspace(δ_μ_min, δ_μ_max, δ_μ_size)  # grid points for δ_μ

    info_matrix = np.empty((ρ_μ_size, δ_μ_size))  # create empty matrix to store stability test results
    GR_matrix = np.empty((ρ_μ_size, δ_μ_size))    # create empty matrix to store G_R values

    for (i_ρ, ρ) in enumerate(ρ_μ_grid):
        for (i_δ, δ) in enumerate(δ_μ_grid):
            ifp_mp = IFP_MeanPersistence(ρ_μ=ρ, δ_μ=δ)
            GR_matrix[i_ρ, i_δ], info_matrix[i_ρ, i_δ] = ifp_mp.stability_test()

    return GR_matrix, info_matrix, ρ_μ_grid, δ_μ_grid

GR_matrix_mp, info_matrix_mp, ρ_μ_grid, δ_μ_grid = stability_range_mp(ρ_μ_size=100, δ_μ_size=100)
```

```python

```

```python
# find the threshold of stability range
idx = np.minimum(np.sum(info_matrix_mp, axis=1), len(δ_μ_grid)-1)
idx = idx.astype(int)       # threshold index on the δ_σ grid points
δ_μ_thresh = δ_μ_grid[idx]  # threshold δ_σ values

fig, ax = plt.subplots(figsize=(8,6))

ax.plot(ρ_μ_grid, δ_μ_thresh, label='stability threshold', lw=2.5)

ax.set_xlabel(r'$\rho_\mu$', size=20)
ax.set_ylabel(r'$\delta_\mu$', size=20)
ax.set_xlim((0.0, 1.0))
ax.set_ylim((0.0, 0.14))
#ax.set_title('Model 2: γ=%i' %γ, size=15)
plt.xticks(fontsize=15)
plt.yticks([0.02,0.04,0.06,0.08,0.10,0.12,0.14], fontsize=15)

font = {'family': 'sans-serif',
        'color':  'black',
        'weight': 'normal',
        'size': 20,
        }

ax.text(0.1, 0.03, 'stable', fontdict=font)
#ax.text(0.3, 0.05, 'unstable', fontdict=font)

ax.plot([0.5722], [0.0067], marker='o', markersize=5, color="blue")
ax.text(0.36, 0.01, "$(0.5722, \, 0.0067)$", family="sans-serif", size=15)

ax.fill_between(ρ_μ_grid, δ_μ_thresh, color='lightyellow')

ax.legend(fontsize=15)

plt.savefig('mp_stb.pdf')

plt.show()
```

```python

```

```python
# ------------------------------------- #
#     make contour plot of G_R
# ------------------------------------- #
fig, ax = plt.subplots(figsize=(10,6))

X2, Y2 = np.meshgrid(ρ_μ_grid, δ_μ_grid)

# hide high values in GR_matrix
Z2 = GR_matrix_mp
Z2[Z2 > 1.5] = np.nan

β = IFP_MeanPersistence().β

CS2_f = ax.contourf(X2, Y2, β*Z2.T, alpha=0.4)
plt.colorbar(CS2_f, ax=ax)

levels = [1.0]

CS2 = ax.contour(X2, Y2, β*Z2.T, colors='k', levels=levels)

ax.clabel(CS2, inline=1, fontsize=15, fmt="%1.2f")

ax.plot([0.5722], [0.0067], marker='o', markersize=5, color="black")
ax.text(0.36, 0.011, "$(0.5722, \, 0.0067)$", family="sans-serif", size=15)

ax.set_xlabel("$\\rho_\mu$", fontsize=20)
ax.set_ylabel("$\delta_\mu$", fontsize=20)

plt.xticks([0.0,0.2,0.4,0.6,0.8,1.0], fontsize=15)
plt.yticks(fontsize=15)

plt.savefig('mp_gr.pdf')

plt.show()
```

```python

```


# 1. Set up the notebook

Import modules.

```python
import numpy as np
import sympy as sym
import json
import matplotlib.pyplot as plt
from scipy import linalg
from scipy.interpolate import interp1d
```

Define a function to load data from a hardware flight test and resample it at 100 Hz using linear interpolation. If `only_in_flight=True`, then only data for which the desired position was positive (i.e., "trying to fly" rather than "sitting on the ground") will be returned.

```python
def load_hardware_data(filename, t_min_offset=0, t_max_offset=0, only_in_flight=False):
    # load raw data
    with open(filename, 'r') as f:
        data = json.load(f)

    # convert lists to numpy arrays
    for val in data.values():
        for key in val.keys():
            val[key] = np.array(val[key])

    # create an array of times at which to subsample
    t_min = -np.inf
    t_max = np.inf
    for key, val in data.items():
        t_min = max(t_min, val['time'][0])
        t_max = min(t_max, val['time'][-1])
    t_min += t_min_offset * 1000
    t_max -= t_max_offset * 1000
    nt = int(1 + np.floor((t_max - t_min) / 10.))
    t = np.arange(0, 10 * nt, 10) / 1000.
    resampled_data = {'time': t}

    # resample raw data with linear interpolation
    for k, v in data.items():
        f = interp1d((v['time'] - t_min) / 1000., v['data'])
        resampled_data[k] = f(t)

    # truncate to times when o_z_des is positive
    if only_in_flight:
        if 'ae483log.o_z_des' not in resampled_data.keys():
            raise Exception('"ae483log.o_z_des" must be logged')
        i = np.argwhere(resampled_data['ae483log.o_z_des'] > 0).flatten()
        if len(i) == 0:
            raise Exception('o_z_des was never positive')
        if len(i) < 2:
            raise Exception('o_z_des was only positive for one time step')
        for key in resampled_data.keys():
            resampled_data[key] = resampled_data[key][i[0]:i[-1]]

    # return the resampled data
    return resampled_data
```

Define a function to solve the linear quadratic regulator (LQR) problem - i.e., that finds the matrix $K$ for which

$$u(t) = - K x(t)$$

is the optimal solution to

$$
\begin{align*}
\underset{u_{[t_{0}, \infty)}}{\text{minimize}} &\qquad \int_{t_{0}}^{\infty}\left( x(t)^{T}Qx(t)+u(t)^{T}Ru(t)\right)dt \\
\text{subject to} &\qquad \dot{x}(t) = Ax(t)+Bu(t), \quad x(t_{0})=x_{0}.
\end{align*}
$$

```python
def lqr(A, B, Q, R):
    P = linalg.solve_continuous_are(A, B, Q, R)
    K = linalg.inv(R) @  B.T @ P
    return K
```

Define a function to create the observability matrix

$$ W_o = \begin{bmatrix} C \\ CA \\ CA^2 \\ \vdots \\ CA^{n-1} \end{bmatrix} $$

where $A$ has size $n \times n$.

```python
def obsv(A, C):
    W = C
    for i in range(1, A.shape[0]):
        W = np.vstack([W, C @ np.linalg.matrix_power(A, i)])
    return W
```

# 2. Design observer

## 2.1 Define state-space model

Here are the matrices that describe the state-space model

$$\begin{align*} \dot{x} &= Ax + Bu \\ y &= Cx + Bu \end{align*}$$

that was derived in Lab 7 (copy/pasted from that week's analysis notebook).

```python
# FIXME (copy/paste definition of A, B, C, D)

```

## 2.2 Show that not all states are observable

Find the observability matrix

$$ W_o = \begin{bmatrix} C \\ CA \\ CA^2 \\ \vdots \\ CA^{n-1} \end{bmatrix} $$

where $A$ is $n \times n$.

```python
W_o = obsv(A, C)
```

Find the rank of the observability matrix using [numpy.linalg.matrix_rank](https://numpy.org/doc/stable/reference/generated/numpy.linalg.matrix_rank.html). The system is **observable** if and only if $W_o$ is **full rank**, that is, if its rank is equal to the number of states $n$.

```python
print(f'      rank of W_o is: {np.linalg.matrix_rank(W_o)}')
print(f'"full rank" would be: {A.shape[0]}')
```

If the system is *not* observable, then it is impossible to design a stable observer - that is, an observer that makes the error in your estimate of each state converge to zero over time.

In particular, the following code would produce an error if you tried to use it:
```python
# Choose weights
Q = np.eye(3) # <-- one diagonal entry for each output
R = np.eye(9) # <-- one diagonal entry for each state

# Find gain matrix
L = lqr(A.T, C.T, linalg.inv(R), linalg.inv(Q)).T
```
It does not matter what method you use - if it is impossible to design a stable observer, that's it, you're out of luck!

## 2.3 Choose a subset of states that are observable

List the index of each state to include.

```python
# FIXME: delete the index of each state that is not observable!

s_obs_index = [
    0, # o_x
    1, # o_y
    2, # o_z
    3, # psi
    4, # theta
    5, # phi
    6, # v_x
    7, # v_y
    8, # v_z
]
```

Define a state-space model

$$
\begin{align*}
\dot{x}_\text{obs} = A_\text{obs} x_\text{obs} + B_\text{obs} u \\
y = C_\text{obs} x_\text{obs} + D_\text{obs} u
\end{align*}
$$

with only these states.

```python
A_obs = A[s_obs_index, :][:, s_obs_index]
B_obs = B[s_obs_index, :]
C_obs = C[:, s_obs_index]
D_obs = D
```

Show the matrices that describe this state-space model.

```python
print(f'A_obs:\n{A_obs}\n')
print(f'B_obs:\n{B_obs}\n')
print(f'C_obs:\n{C_obs}\n')
print(f'D_obs:\n{D_obs}\n')
```

Check that this subsystem is observable:

```python
print(f'      rank of W_o is: {np.linalg.matrix_rank(obsv(A_obs, C_obs))}')
print(f'"full rank" would be: {A_obs.shape[0]}')
```

## 2.4 Choose gain matrix for the observable subsystem

### 2.4.1 With equal weights

Choose weights $Q$ and $R$ as identity matrices.

```python
# FIXME: change the size of each identity matrix to match the
# number of states and outputs for your observable subsystem

Q = np.eye(3) # <-- one diagonal entry for each output
R = np.eye(9) # <-- one diagonal entry for each state
```

Find gain matrix $L$ for the chosen $Q$ and $R$ by solving an LQR problem.

```python
L = lqr(A_obs.T, C_obs.T, linalg.inv(R), linalg.inv(Q)).T
```

Show $L$ (formatted nicely).

```python
L_str = np.array2string(L,
                        formatter={'float_kind': lambda x: f'{x:12.6f}'},
                        prefix='    ',
                        max_line_width=np.inf)

print(f'L = {L_str}')
```

### 2.4.2 With weights chosen by model error variance

Choose weights $Q$ and $R$ based on variance of model error.

```python
# FIXME:
#
# (1) Change the size of each matrix to match the number of
#     states and outputs for your observable subsystem
#
# (2) Change the value of each diagonal entry based on your
#     results from Lab 7.

Q = np.diag([
    1.,             # n_x
    1.,             # n_y
    1.,             # r
])

R = np.diag([
    1., # o_x
    1., # o_y
    1., # o_z
    1., # psi
    1., # theta
    1., # phi
    1., # v_x
    1., # v_y
    1., # v_z
])
```

Find gain matrix $L$ for the chosen $Q$ and $R$ by solving an LQR problem.

```python
L = lqr(A_obs.T, C_obs.T, linalg.inv(R), linalg.inv(Q)).T
```

Show $L$ (formatted nicely).

```python
L_str = np.array2string(L,
                        formatter={'float_kind': lambda x: f'{x:12.6f}'},
                        prefix='    ',
                        max_line_width=np.inf)

print(f'L = {L_str}')
```

# 3. Implement and test observer (offline)

## 3.1 Do a flight test

#### Describe the flight test

Replace this cell with the following information, at minimum:

* A description of the flight trajectory, both in words and (if you like) a snippet of code from `flight.py`.
* A description of the flight conditions (e.g., where was the flight conducted, did you power cycle the drone just before flying, were you using the positioning system or only the onboard sensors, etc.).

Also fill in the following table, which describes the point at which the drone landed, as measured with a ruler:

| | $o_x$ (meters) | $o_y$ (meters) |
| ---: | :---: | :---: |
| landing position | 1.000 | 2.000 |

#### Show a video of the flight test

Replace this cell with a video of your flight. Here are two ways to do this. (What you should *not* do is drag-and-drop your video into this notebook.)

##### Markdown

Put your video in the same directory as this notebook. Suppose this video is called `hardware_video.mov`. Then put the following code in a cell of type `Markdown` and evaluate it:
```

```

##### HTML
Put your video in the same directory as this notebook. Suppose this video is called `hardware_video.mov`. Then put the following code in a cell of type `Code` and evaluate it:
```
%%HTML

```
You can change the `width` parameter to resize your video.

## 3.2 Load and parse data from flight test

Load flight test data.

```python
data = load_hardware_data(
    'hardware_data.json',   # <-- FIXME: replace with name of file with data
    t_min_offset=0.,
    t_max_offset=0.,
    only_in_flight=True,
)
```

Parse flight test data.

```python
# time
t = data['time']

# states
o_x_true = data['ae483log.o_x']
o_y_true = data['ae483log.o_y']
o_z_true = data['ae483log.o_z']
psi_true = data['ae483log.psi']
theta_true = data['ae483log.theta']
phi_true = data['ae483log.phi']
v_x_true = data['ae483log.v_x']
v_y_true = data['ae483log.v_y']
v_z_true = data['ae483log.v_z']

# inputs
w_x_true = data['ae483log.w_x']
w_y_true = data['ae483log.w_y']
w_z_true = data['ae483log.w_z']
a_z_true = data['ae483log.a_z']

# outputs
n_x_true = data['ae483log.n_x']
n_y_true = data['ae483log.n_y']
r_true = data['ae483log.r']
```

Define parameter values.

```python
# Acceleration of gravity
g = 9.81

# Optical flow constant (do not change)
k_flow = 0.01 * 30.0 / np.deg2rad(4.2)

# Equilibrium value of o_z
o_z_eq = 0. # <-- FIXME: replace with the equilibrium height you chose in Lab 7

# Time step (should be 0.01)
dt = t[1] - t[0]
print(f'dt = {dt:.4f} (should be 0.01)')
```

## 3.3 Apply observer for gain matrix chosen by equal weights

Implement and run your observer in the following code cell. For each state, the result will be a numpy array (e.g., `o_x_hat`) that contains a state estimate at every time step.

```python
# Initialize estimates
o_x = 0.
o_y = 0.
o_z = 0.
psi = 0.
theta = 0.
phi = 0.
v_x = 0.
v_y = 0.
v_z = 0.

# Initialize storage (do not change)
o_x_hat = np.zeros(len(t))
o_y_hat = np.zeros(len(t))
o_z_hat = np.zeros(len(t))
psi_hat = np.zeros(len(t))
theta_hat = np.zeros(len(t))
phi_hat = np.zeros(len(t))
v_x_hat = np.zeros(len(t))
v_y_hat = np.zeros(len(t))
v_z_hat = np.zeros(len(t))

# Iterate over all time steps
for i in range(len(t)):
    # Get measurements (do not change)
    w_x = w_x_true[i]
    w_y = w_y_true[i]
    w_z = w_z_true[i]
    a_z = a_z_true[i]
    n_x = n_x_true[i]
    n_y = n_y_true[i]
    r = r_true[i]

    # Compute each element of:
    #
    #  C x + D u - y
    #
    n_x_err = 0.  # <-- FIXME
    n_y_err = 0.  # <-- FIXME
    r_err = 0.    # <-- FIXME

    # Update estimates
    o_x += dt * (0.)   # <-- FIXME
    o_y += dt * (0.)   # <-- FIXME
    o_z += dt * (0.)   # <-- FIXME
    psi += dt * (0.)   # <-- FIXME
    theta += dt * (0.) # <-- FIXME
    phi += dt * (0.)   # <-- FIXME
    v_x += dt * (0.)   # <-- FIXME
    v_y += dt * (0.)   # <-- FIXME
    v_z += dt * (0.)   # <-- FIXME

    # Store estimates (do not change)
    o_x_hat[i] = o_x
    o_y_hat[i] = o_y
    o_z_hat[i] = o_z
    psi_hat[i] = psi
    theta_hat[i] = theta
    phi_hat[i] = phi
    v_x_hat[i] = v_x
    v_y_hat[i] = v_y
    v_z_hat[i] = v_z
```

Compute error in each state estimate (relative to "ground truth" from the stock code).

```python
o_x_err = o_x_hat - o_x_true
o_y_err = o_y_hat - o_y_true
o_z_err = o_z_hat - o_z_true
psi_err = psi_hat - psi_true
theta_err = theta_hat - theta_true
phi_err = phi_hat - phi_true
v_x_err = v_x_hat - v_x_true
v_y_err = v_y_hat - v_y_true
v_z_err = v_z_hat - v_z_true
```

Compare estimates with ground truth.

```python
def comparison_plot(t, s_hat, s_true, name, ax):
    ax.plot(t, s_hat, label=f'{name} (estimated)', linewidth=3)
    ax.plot(t, s_true, '--', label=f'{name} (true)', linewidth=2)
    ax.legend()
    ax.grid()

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25), sharex=True)
comparison_plot(t, o_x_hat, o_x_true, 'o_x', ax_o_x)
comparison_plot(t, o_y_hat, o_y_true, 'o_y', ax_o_y)
comparison_plot(t, o_z_hat, o_z_true, 'o_z', ax_o_z)
comparison_plot(t, psi_hat, psi_true, 'psi', ax_psi)
comparison_plot(t, theta_hat, theta_true, 'theta', ax_theta)
comparison_plot(t, phi_hat, phi_true, 'phi', ax_phi)
comparison_plot(t, v_x_hat, v_x_true, 'v_x', ax_v_x)
comparison_plot(t, v_y_hat, v_y_true, 'v_y', ax_v_y)
comparison_plot(t, v_z_hat, v_z_true, 'v_z', ax_v_z)
ax_v_z.set_xlabel('time (s)')
plt.show()
```

Plot error in each state estimate with respect to ground truth.

```python
def error_plot(t, s_err, name, ax):
    ax.plot(t, s_err, label=f'{name} (error)', linewidth=3)
    ax.legend()
    ax.grid()

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25), sharex=True)
error_plot(t, o_x_err, 'o_x', ax_o_x)
error_plot(t, o_y_err, 'o_y', ax_o_y)
error_plot(t, o_z_err, 'o_z', ax_o_z)
error_plot(t, psi_err, 'psi', ax_psi)
error_plot(t, theta_err, 'theta', ax_theta)
error_plot(t, phi_err, 'phi', ax_phi)
error_plot(t, v_x_err, 'v_x', ax_v_x)
error_plot(t, v_y_err, 'v_y', ax_v_y)
error_plot(t, v_z_err, 'v_z', ax_v_z)
ax_v_z.set_xlabel('time (s)')
plt.show()
```

Plot histogram of errors in each state estimate.

```python
def histogram_plot(t, s_err, name, ax):
    label = f'Error in estimate of {name} ' + \
            f'(RMSE = {np.sqrt(np.mean(s_err**2)):6.3f}, ' + \
            f'mean = {np.mean(s_err):6.3f}, ' + \
            f'std = {np.std(s_err):6.3f})'
    ax.hist(s_err, 50, label=label)
    ax.legend(fontsize=14)

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25))
histogram_plot(t, o_x_err, 'o_x', ax_o_x)
histogram_plot(t, o_y_err, 'o_y', ax_o_y)
histogram_plot(t, o_z_err, 'o_z', ax_o_z)
histogram_plot(t, psi_err, 'psi', ax_psi)
histogram_plot(t, theta_err, 'theta', ax_theta)
histogram_plot(t, phi_err, 'phi', ax_phi)
histogram_plot(t, v_x_err, 'v_x', ax_v_x)
histogram_plot(t, v_y_err, 'v_y', ax_v_y)
histogram_plot(t, v_z_err, 'v_z', ax_v_z)
fig.tight_layout()
plt.show()
```

## 3.4 Apply observer for gain matrix chosen by weights from model error variance

Implement and run your observer in the following code cell. For each state, the result will be a numpy array (e.g., `o_x_hat`) that contains a state estimate at every time step.

```python
# Initialize estimates
o_x = 0.
o_y = 0.
o_z = 0.
psi = 0.
theta = 0.
phi = 0.
v_x = 0.
v_y = 0.
v_z = 0.

# Initialize storage (do not change)
o_x_hat = np.zeros(len(t))
o_y_hat = np.zeros(len(t))
o_z_hat = np.zeros(len(t))
psi_hat = np.zeros(len(t))
theta_hat = np.zeros(len(t))
phi_hat = np.zeros(len(t))
v_x_hat = np.zeros(len(t))
v_y_hat = np.zeros(len(t))
v_z_hat = np.zeros(len(t))

# Iterate over all time steps
for i in range(len(t)):
    # Get measurements (do not change)
    w_x = w_x_true[i]
    w_y = w_y_true[i]
    w_z = w_z_true[i]
    a_z = a_z_true[i]
    n_x = n_x_true[i]
    n_y = n_y_true[i]
    r = r_true[i]

    # Compute each element of:
    #
    #  C x + D u - y
    #
    n_x_err = 0.  # <-- FIXME
    n_y_err = 0.  # <-- FIXME
    r_err = 0.    # <-- FIXME

    # Update estimates
    o_x += dt * (0.)   # <-- FIXME
    o_y += dt * (0.)   # <-- FIXME
    o_z += dt * (0.)   # <-- FIXME
    psi += dt * (0.)   # <-- FIXME
    theta += dt * (0.) # <-- FIXME
    phi += dt * (0.)   # <-- FIXME
    v_x += dt * (0.)   # <-- FIXME
    v_y += dt * (0.)   # <-- FIXME
    v_z += dt * (0.)   # <-- FIXME

    # Store estimates (do not change)
    o_x_hat[i] = o_x
    o_y_hat[i] = o_y
    o_z_hat[i] = o_z
    psi_hat[i] = psi
    theta_hat[i] = theta
    phi_hat[i] = phi
    v_x_hat[i] = v_x
    v_y_hat[i] = v_y
    v_z_hat[i] = v_z
```

Compute error in each state estimate (relative to "ground truth" from the stock code).

```python
o_x_err = o_x_hat - o_x_true
o_y_err = o_y_hat - o_y_true
o_z_err = o_z_hat - o_z_true
psi_err = psi_hat - psi_true
theta_err = theta_hat - theta_true
phi_err = phi_hat - phi_true
v_x_err = v_x_hat - v_x_true
v_y_err = v_y_hat - v_y_true
v_z_err = v_z_hat - v_z_true
```

Compare estimates with ground truth.

```python
def comparison_plot(t, s_hat, s_true, name, ax):
    ax.plot(t, s_hat, label=f'{name} (estimated)', linewidth=3)
    ax.plot(t, s_true, '--', label=f'{name} (true)', linewidth=2)
    ax.legend()
    ax.grid()

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25), sharex=True)
comparison_plot(t, o_x_hat, o_x_true, 'o_x', ax_o_x)
comparison_plot(t, o_y_hat, o_y_true, 'o_y', ax_o_y)
comparison_plot(t, o_z_hat, o_z_true, 'o_z', ax_o_z)
comparison_plot(t, psi_hat, psi_true, 'psi', ax_psi)
comparison_plot(t, theta_hat, theta_true, 'theta', ax_theta)
comparison_plot(t, phi_hat, phi_true, 'phi', ax_phi)
comparison_plot(t, v_x_hat, v_x_true, 'v_x', ax_v_x)
comparison_plot(t, v_y_hat, v_y_true, 'v_y', ax_v_y)
comparison_plot(t, v_z_hat, v_z_true, 'v_z', ax_v_z)
ax_v_z.set_xlabel('time (s)')
plt.show()
```

Plot error in each state estimate with respect to ground truth.

```python
def error_plot(t, s_err, name, ax):
    ax.plot(t, s_err, label=f'{name} (error)', linewidth=3)
    ax.legend()
    ax.grid()

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25), sharex=True)
error_plot(t, o_x_err, 'o_x', ax_o_x)
error_plot(t, o_y_err, 'o_y', ax_o_y)
error_plot(t, o_z_err, 'o_z', ax_o_z)
error_plot(t, psi_err, 'psi', ax_psi)
error_plot(t, theta_err, 'theta', ax_theta)
error_plot(t, phi_err, 'phi', ax_phi)
error_plot(t, v_x_err, 'v_x', ax_v_x)
error_plot(t, v_y_err, 'v_y', ax_v_y)
error_plot(t, v_z_err, 'v_z', ax_v_z)
ax_v_z.set_xlabel('time (s)')
plt.show()
```

Plot histogram of errors in each state estimate.

```python
def histogram_plot(t, s_err, name, ax):
    label = f'Error in estimate of {name} ' + \
            f'(RMSE = {np.sqrt(np.mean(s_err**2)):6.3f}, ' + \
            f'mean = {np.mean(s_err):6.3f}, ' + \
            f'std = {np.std(s_err):6.3f})'
    ax.hist(s_err, 50, label=label)
    ax.legend(fontsize=14)

fig, (ax_o_x,
      ax_o_y,
      ax_o_z,
      ax_psi,
      ax_theta,
      ax_phi,
      ax_v_x,
      ax_v_y,
      ax_v_z) = plt.subplots(9, 1, figsize=(15, 25))
histogram_plot(t, o_x_err, 'o_x', ax_o_x)
histogram_plot(t, o_y_err, 'o_y', ax_o_y)
histogram_plot(t, o_z_err, 'o_z', ax_o_z)
histogram_plot(t, psi_err, 'psi', ax_psi)
histogram_plot(t, theta_err, 'theta', ax_theta)
histogram_plot(t, phi_err, 'phi', ax_phi)
histogram_plot(t, v_x_err, 'v_x', ax_v_x)
histogram_plot(t, v_y_err, 'v_y', ax_v_y)
histogram_plot(t, v_z_err, 'v_z', ax_v_z)
fig.tight_layout()
plt.show()
```

## 3.5 Summary and discussion

The following table reports the RMSE in each state estimate:

|  | $o_x$ | $o_y$ | $o_z$ | $\psi$ | $\theta$ | $\phi$ | $v_x$ | $v_y$ | $v_z$ |
| --: | :--: | :--: | :--: | :--: | :--: | :--: | :--: | :--: | :--: |
| RMSE with equal weights | 0.001 | 0.002 | 0.003 | 0.004 | 0.005 | 0.006 | 0.007 | 0.008 | 0.009 |
| RMSE with weights from error model variance | 0.001 | 0.002 | 0.003 | 0.004 | 0.005 | 0.006 | 0.007 | 0.008 | 0.009 |

**Modify the text in this cell** to answer the following questions:

* Which observer performs better, the one derived from equal weights or the one derive from weights based on error model variance? Why? (Make reference to your table of results when answering this question.)
* Is there a difference in what happens to the error in your estimate of (1) observable states and (2) non-observable states over time? Why? (Make reference to your plots when answering this question.)
* We refer to the state estimates that are computed on-board by the stock code as "true"? Are they? For example, compare the landing position that you measured with a ruler (Section 3.1) to the final position that was estimated by the stock code - to what extent are they the same? (You may want to disinguish between observable and non-observable states when answering this question.)


# 10 Ordinary Differential Equations (ODEs)

[ODE](http://mathworld.wolfram.com/OrdinaryDifferentialEquation.html)s describe many phenomena in physics. They describe the changes of a **dependent variable** $y(t)$ as a function of a **single independent variable** (e.g. $t$ or $x$).

An ODE of **order** $n$

$$
F(t, y^{(0)}, y^{(1)}, ..., y^{(n)}) = 0
$$

contains derivatives $y^{(k)}(t) \equiv y^{(k)} \equiv \frac{d^{k}y(t)}{dt^{k}}$ up to the $n$-th derivative (and $y^{(0)} \equiv y$).

### Initial and boundary conditions

* $n$ **initial conditions** are needed to *uniquely determine* the solution of a $n$-th order ODE, e.g, initial position and velocities.

* **Boundary conditions** (values of solution on domain boundries) can additionaly restrict solutions but the resulting *eigenvalue problems* are more difficult, e.g, wavefunction goes towards 0 for $\pm\infty$.

### Linear ODEs
A **linear** ODE contains no higher powers than 1 of any of the $y^{(k)}$.

*Superposition principle*: Linear combinations of solutions are also solutions.

#### Example: First order linear ODE

\begin{align}
\frac{dy}{dt} &= f(t)y + g(t)\\
y^{(1)} &= f(t)y + g(t)\\
% y^{(1)} - f(t)y - g(t) &= 0
\end{align}

##### Radioactive decay
$$
\frac{dN}{dt} = -k N
$$

### Non-linear ODEs
**Non-linear** ODEs can contain any powers in the dependent variable and its derivatives.

No superposition of solutions. Often impossible to solve analytically.

#### Example: Second order (general) ODE
\begin{gather}
\frac{d^2 y}{dt^2} + \lambda(t) \frac{dy}{dt} = f\left(t, y, \frac{dy}{dt}\right)\\
\end{gather}

##### Newton's equations of motion
$$
m\frac{d^2 x}{dt^2} = F(x) + F_\text{ext}(x, t) \quad \text{with}
\quad F(x) = -\frac{dU}{dx}
$$
(Force is often derived from a potential energy $U(x)$ and may contain non-linear terms such as $x^{-2}$ or $x^3$.)

## Partial differential equations (PDEs)
* more than one independent variable (e.g. $x$ and $t$)
* partial derivatives
* much more difficult than ODEs

#### Example: Schrödinger equation (Quantum Mechanics)
$$
i\hbar \frac{\partial\psi(\mathbf{x}, t)}{\partial t} = -\frac{\hbar^2}{2m}
   \left(\frac{\partial^2 \psi}{\partial x^2} +
         \frac{\partial^2 \psi}{\partial y^2} +
         \frac{\partial^2 \psi}{\partial z^2}
   \right) + V(\mathbf{x})\, \psi(\mathbf{x}, t)
$$

## Harmonic and anharmonic oscillator
* particle with mass $m$ connected to a spring
* spring described by a harmonic potential or anharmonic ones in the displacements from equilibrium $x$

\begin{align}
U_1(x) &= \frac{1}{2} k x^2, \quad k=1\\
U_2(x) &= \frac{1}{2} k x^2 \left(1 - \frac{2}{3}\alpha x\right), \quad k=1,\ \alpha=\frac{1}{2}\\
U_3(x) &= \frac{1}{p} k x^p, \quad k=1,\ p=6
\end{align}

1. What do these potentials look like? Sketch or plot.
2. Calculate the forces.

#### Potentials

```python
import numpy as np

def U1(x, k=1):
    return 0.5 * k * x*x

def U2(x, k=1, alpha=0.5):
    return 0.5 * k * x*x * (1 - (2/3)*alpha*x)

def U3(x, k=1, p=6):
    return (k/p) * np.power(x, p)
```

```python
import matplotlib
import matplotlib.pyplot as plt
matplotlib.style.use('seaborn-talk')
%matplotlib inline
```

```python
X = np.linspace(-3, 3, 100)
ax = plt.subplot(1,1,1)
ax.plot(X, U1(X), label=r"$U_1$")
ax.plot(X, U2(X), label=r"$U_2$")
ax.plot(X, U3(X), label=r"$U_3$")
ax.set_ylim(-0.5, 10)
ax.legend(loc="upper center");
```

#### Forces

\begin{align}
F_1(x) &= -kx\\
F_2(x) &= -kx(1 + \alpha x)\\
F_3(x) &= -k x^{p-1}
\end{align}

## ODE Algorithms
Basic idea:

1. Start with initial conditions, $y_0 \equiv y(t=0)$
2. Use $\frac{dy}{dt} = f(t, y)$ (the RHS!) to advance solution a small step $h$ forward in time: $y(t=h) \equiv y_1$
3. Repeat with $y_1$ to obtain $y_2 \equiv y(t=2h)$... and for all future values of $t$.

Possible issues
* small differences: subtractive cancelation and round-off error accumulation
* extrapolation: numerical "solution" can deviate wildly from exact
* possibly need adaptive $h$

### Euler's rule
Simple: forward difference

\begin{align}
f(t, y) = \frac{dy(t)}{dt} &\approx \frac{y(t_{n+1}) - y(t_n)}{h}\\
y_{n+1} &\approx y_n + h f(t_n, y_n) \quad \text{with} \quad y_n := y(t_n)
\end{align}

Error will be $\mathcal{O}(h^2)$ (bad!).

Also: what if we have a second order ODE ?!?! We only used $dy/dt$.

### Convert 2nd order ODE to 2 coupled _First Order_ ODEs

### Convert 2nd order ODE to 2 coupled 1st order ODEs
The 2nd order ODE is
$$
\frac{d^2 y}{dt^2} = f(t, y)
$$

Introduce "dummy" dependent variables $y_i$ with $y_0 \equiv y$ and

\begin{alignat}{1}
\frac{dy}{dt} &= \frac{dy_0}{dt} &=   y_1\\
\frac{d^2y}{dt^2} &= \frac{dy_1}{dt} &= {} f(t, y_0).
\end{alignat}

The first equation defines the velocity $y_1 = v$ and the second one is the original ODE.

### $n$-th oder ODE to $n$ coupled 1st order ODEs

The $n$-th order ODE is
$$
\frac{d^n y}{dt^n} = f(t, y, \frac{d y}{dt}, \frac{d^2 y}{dt^2}, \dots, \frac{d^{n-1} y}{dt^{n-1}})
$$

Introduce "dummy" dependent variables $y^{(i)}$ with $y^{(0)} \equiv y$ and

\begin{align}
\frac{dy^{(0)}}{dt} &=  y^{(1)}\\
\frac{dy^{(1)}}{dt} &=  y^{(2)}\\
\dots & \\
\frac{dy^{(n-1)}}{dt} &=  f(t, y^{(0)}, y^{(1)}, y^{(2)}, \dots, y^{(n-1)}).
\end{align}

### General standard (dynamic) form of ODEs
1 ODE of *any order* $n$ $\rightarrow$ $n$ coupled simultaneous first-order ODEs in $n$ unknowns  $y^{(0)}, \dots, y^{(n-1)}$:

\begin{align}
\frac{dy^{(0)}}{dt} &= f^{(0)}(t, y^{(0)}, \dots, y^{(n-1)})\\
\frac{dy^{(1)}}{dt} &= f^{(1)}(t, y^{(0)}, \dots, y^{(n-1)})\\
\vdots &  \\
\frac{dy^{(n-1)}}{dt} &= f^{(n-1)}(t, y^{(0)}, \dots, y^{(n-1)})\\
\end{align}

In $n$-dimensional vector notation:

\begin{align}
\frac{d\mathbf{y}(t)}{dt} &= \mathbf{f}(t, \mathbf{y})\\
\mathbf{y} &= \left(\begin{array}{c}
  y^{(0)}(t) \\
  y^{(1)}(t) \\
  \vdots \\
  y^{(n-1)}(t)
  \end{array}\right),
\quad
\mathbf{f} = \left(\begin{array}{c}
  f^{(0)}(t, \mathbf{y}) \\
  f^{(1)}(t, \mathbf{y}) \\
  \vdots \\
  f^{(n-1)}(t, \mathbf{y})
  \end{array}\right)
\end{align}

#### Example: Convert Newton's EOMs to standard form

$$
\frac{d^2 x}{dt^2} = m^{-1} F\Big(t, x, \frac{dx}{dt}\Big)
$$

RHS may *not contain any explicit derivatives* but components of $\mathbf{y}$ can represent derivatives.
* position $x$ as first dependent variable $y^{(0)}$ (as usual).
* velocity $dx/dt$ as second dependent variable $y^{(1)}$

\begin{align}
y^{(0)}(t) &:= x(t)\\
y^{(1)}(t) &:= \frac{dx}{dt} = \frac{dy^{(0)}}{dt}
\end{align}

One 2nd order ODE

$$
\frac{d^2 x}{dt^2} = m^{-1} F\Big(t, x, \frac{dx}{dt}\Big)
$$

to two simultaneous 1st order ODEs:

\begin{align}
\frac{dy^{(0)}}{dt} &= y^{(1)}(t)\\
\frac{dy^{(1)}}{dt} &= m^{-1} F\Big(t, y^{(0)}, y^{(1)}\Big)
\end{align}

\begin{align}
\frac{d\mathbf{y}(t)}{dt} &= \mathbf{f}(t, \mathbf{y})\\
\mathbf{y} &= \left(\begin{array}{c}
     y^{(0)} \\
     y^{(1)}
  \end{array}\right) =
  \left(\begin{array}{c}
     x(t) \\
     \frac{dx}{dt}
  \end{array}\right),\\
\mathbf{f} &= \left(\begin{array}{c}
    y^{(1)}(t) \\
    m^{-1} F\Big(t, y^{(0)}, y^{(1)}\Big)
  \end{array}\right) =
  \left(\begin{array}{c}
    \frac{dx}{dt} \\
    m^{-1} F\Big(t, x(t), \frac{dx}{dt}\Big)
  \end{array}\right)
\end{align}

#### Example:  1D harmonic oscillator in standard form
With $F_1 = -k x$:
$$
\frac{d^2 x}{dt^2} = -m^{-1}k x
$$
convert to

\begin{align}
\frac{dy^{(0)}}{dt} &= y^{(1)}(t) \\
\frac{dy^{(1)}}{dt} &= -m^{-1}k y^{(0)}
\end{align}

Force (or derivative) function $\mathbf{f}$ and initial conditions:

\begin{alignat}{3}
f^{(0)}(t, \mathbf{y}) &= y^{(1)},
    &\quad y^{(0)}(0) &= x_0,\\
f^{(1)}(t, \mathbf{y}) &=  -m^{-1} k y^{(0)},
    &\quad y^{(1)}(0) &= v_0.
\end{alignat}

### Euler's rule (standard form)
Given the $n$-dimensional vectors from the ODE standard form

$$
\frac{d\mathbf{y}}{dt} = \mathbf{f}(t, \mathbf{y})
$$
the **Euler rule** amounts to

\begin{align}
\mathbf{f}(t, \mathbf{y}) = \frac{d\mathbf{y}(t)}{dt} &\approx \frac{\mathbf{y}(t_{n+1}) - \mathbf{y}(t_n)}{\Delta t}\\
\mathbf{y}_{n+1} &\approx \mathbf{y}_n + \Delta t \mathbf{f}(t_n, \mathbf{y}_n) \quad \text{with} \quad \mathbf{y}_n := \mathbf{y}(t_n)
\end{align}

## Problem: Numerically integrate the 1D harmonic oscillator with Euler

\begin{alignat}{3}
f^{(0)}(t, \mathbf{y}) &= y^{(1)},
    &\quad y^{(0)}(0) &= x_0,\\
f^{(1)}(t, \mathbf{y}) &=   - \frac{k}{m} y^{(0)},
    &\quad y^{(1)}(0) &= v_0.
\end{alignat}

with $k=1$; $x_0 = 0$ and $v_0 = +1$.

### Explicit implementation:
* Note how in `f_harmonic` we are constructing the force vector of the standard ODE representation
* `y` is the vector of dependents in the standard representation
* We pre-allocate the array for `y` and then assign to individual elements with the
  ```python
  y[:] = ...
  ```
  notation, which has higher performance than creating the array anew every time.

```python
import numpy as np

def F1(x, k=1):
    """Harmonic force"""
    return -k*x

def f_harmonic(t, y, k=1, m=1):
    """Force vector in standard ODE form (n=2)"""
    return np.array([y[1], F1(y[0], k=k)/m])

t_max = 100
h = 0.01
Nsteps = int(t_max/h)
t_range = h * np.arange(Nsteps)
x = np.empty_like(t_range)

y = np.zeros(2)

# initial conditions
x0, v0 = 0.0, 1.0
y[:] = x0, v0

for i, t in enumerate(t_range):
    # store position that corresponds to time t_i
    x[i] = y[0]
    # Euler integrator
    y[:] = y + h * f_harmonic(t, y)
```

Plot the position $x(t)$ (which is $y_0$) against time:

```python
plt.plot(t_range, x)
```

Although we see oscillations in $x(t)$, the fact that the *amplitude increases with time* instead of staying staying constant (as required by *energy conservation*) indicates that the Euler integrator has a serious problem. We will come back to this point later.

### Modular solution with functions
We can make the Euler integrator a function, which makes the code more readable and modular and we can make the whole integration a function, too. This will allow us to easily run the integration with different initial values or `h` steps.

```python
import numpy as np

def F1(x, k=1):
    """Harmonic force"""
    return -k*x

def f_harmonic(t, y, k=1, m=1):
    """Force vector in standard ODE form (n=2)"""
    return np.array([y[1], F1(y[0], k=k)/m])

def euler(y, f, t, h):
    """Euler integrator.

    Returns new y at t+h.
    """
    return y + h * f(t, y)

def integrate(x0=0, v0=1, t_max=100, h=0.001):
    """Integrate the harmonic oscillator with force F1.

    Note that the spring constant k and particle mass m are currently
    pre-defined.

    Arguments
    ---------
    x0 : float
       initial position
    v0 : float
       initial velocity
    t_max : float
       time to integrate out to
    h : float, default 0.001
       integration time step

    Returns
    -------
    Tuple ``(t, x)`` with times and positions.
    """

    Nsteps = t_max/h
    t_range = h * np.arange(Nsteps)
    x = np.empty_like(t_range)
    y = np.zeros(2)
    # initial conditions
    y[:] = x0, v0

    for i, t in enumerate(t_range):
        # store position that corresponds to time t_i
        x[i] = y[0]
        # Euler integrator
        y[:] = euler(y, f_harmonic, t, h)

    return t_range, x
```

Plot the position as a function of time, $x(t)$.

```python
t, x = integrate(t_max=300, h=0.001)
plt.plot(t, x)
```

### Compare to analytical solution

Analytical solution:
$$
\omega = \sqrt{\frac{k}{m}}
$$
and
$$
x(t) = A \sin(\omega t)
$$
here with $A=1$ and $\omega = 1$:

```python
plt.plot(t, x)
plt.plot(t, np.sin(t), color="red")
```

Note the increase in amplitude. Explore if smaller $h$ fixes this obvious problem.

```python
t, x = integrate(h=0.001)
plt.plot(t, x)
plt.plot(t, np.sin(t), color="red")
```

Smaller $h$ improves the integration (but Euler is still a bad algorithm... just run out for longer, i.e., higher `t_max`.)

### Note on Euler's global error

How does the global error after time $t$ behave, i.e., what is its dependency on step size $h$?

* Per-step (**local**) error is $\mathcal{O}(h^2)$.
* For time $t = N h$ we have $N \propto 1/h$ steps.

We assume that we linearly accumulate local error for the **global error** for $N$ steps:

$$
\mathcal{O}(h^2) \times h^{-1} \propto \mathcal{O}(h)
$$

Thus, halving $h$ should halve the global error.

## Optional: Semi-implicit Euler

We can apply a simple trick to turn standard (forward) Euler into a _much_ better algorithm but this only works nicely when we solve Hamilton's equations of motion (Newton's equation without velocity-dependent potentials) and it also does not work nicely with the standard form.

In the following, $x$ is the position, $v = \dot{x} = \frac{dx}{dt}$ the velocity, and $f = F/m$ is the acceleration (but keeping the suggestive letter "f" for the force). Note that the force only depends on positions, $f(x) = F(x)/m$.

### Standard Euler for Hamiltonian dynamics

\begin{align}
x_{n+1} &= x_n + h v_n\\
v_{n+1} &= v_n + h f(x_n)
\end{align}

This is the same algorithm that we discussed in a more general form above.

### Semi-implicit Euler for Hamiltonian dynamics

The trick is to use the _updated_ positions $x_{n+1}$ for the velocity $v_{n+1}$ update, which gives the [semi-implicit Euler](https://en.wikipedia.org/wiki/Semi-implicit_Euler_method) method:

\begin{align}
x_{n+1} &= x_n + h v_n\\
v_{n+1} &= v_n + h f(x_{n+1})
\end{align}

This algorithm _does_ conserve energy over the long term because it is a symplectic algorithm (it maintains the deeper mathematical structures and invariants of Hamiltonian dynamics) although its error is still the same as for standard Euler.

```python

```


```python
%load_ext autoreload
from mps.state import *
```

# TIME EVOLUTION

```python
# file: mps/evolution.py

import numpy as np
import scipy.linalg
from numbers import Number
import mps.state
import scipy.sparse as sp
from mps.state import _truncate_vector, DEFAULT_TOLERANCE
```

## Suzuki-Trotter Decomposition

In Suzuki-Trotter decomposition, the Hamiltonians of the nearest neighbor couplings can be decomposed into two non-commuting parts, $H_{\text{odd}} $ and $ H_{\text{even}} $, so that all additive 2-site operators in each part commute with each other.

Let us consider a simple example of tight binding model with on-site potential and decompose the Hamiltonian into 2-site terms, so that $H=\sum_i h_{i,i+1}$.
\begin{equation}
h_{i,i+1} = \left(\frac{\omega}{2}  a_i^\dagger a_i \right) + \left(\frac{\omega}{2}  a_{i+1}^\dagger a_{i+1} \right) - \left( t a_{i}^\dagger a_{i+1} + \text{h.c.} \right).
\end{equation}
Since $[h_{i,i+1},h_{i+2,i+3}] = 0$, we can group these terms for even and odd $i$, so that $H = H_{\text{odd}} + H_{\text{even}} $.

Note that the local term $ a_i^\dagger a_i$ appears only in one of the groups for $i=1$ and $i=N$. Therefore we need to add two on-site terms $h_1 = \left(\frac{\omega}{2}  a_1^\dagger a_1 \right) $ and $h_N = \left(\frac{\omega}{2}  a_N^\dagger a_N \right) $, to the corresponding two-site terms. So that $h_{1,2} \rightarrow h_{1,2} + h_1$, and $h_{N-1,N} \rightarrow h_{N-1,N} + h_N$.

And for the first order Suzuki-Trotter decomposition, the evolution operator becomes
\begin{equation}
e^{-i \hat{H} \Delta t} = e^{-i \hat{H}_{\text{odd}} \Delta t}  e^{-i \hat{H}_{\text{even}} \Delta t} + O(\Delta t^2).
\end{equation}

`pairwise_unitaries` creates a list of Trotter unitarities corresponding to two-site operators, $U_{i,i+1} = e^{-i h_{i,i+1} \Delta t}$. The Trotter unitarities associated with $\hat{H}_{\text{odd}}$ and $\hat{H}_{\text{even}}$ are applied separately in consecutive sweeps depending on evenodd value passed to TEBD_sweep class:
$$ U = [U_{1,2}, U_{2,3}, U_{3,4}, \dots ]. $$

```python
# file: mps/evolution.py

def pairwise_unitaries(H, δt):
    return [scipy.linalg.expm((-1j * δt) * H.interaction_term(k)).
                              reshape(H.dimension(k), H.dimension(k+1),
                                      H.dimension(k), H.dimension(k+1))
            for k in range(H.size-1)]
```

We apply each $U_{i,i+1} = e^{-i h_{i,i+1} \Delta t}$ to two neighbouring tensors, $A_i$ and $A_{i+1}$ simultaneously, as shown below.

The resulting tensor $B$ is a two-site tensor. We split this tensor using the canonical form algorithm defined in [this notebook](File%201c%20-%20Canonical%20form.ipynb).

```python
# file: mps/evolution.py

def apply_pairwise_unitaries(U, ψ, start, direction, tol=DEFAULT_TOLERANCE):
    """Apply the list of pairwise unitaries U onto an MPS state ψ in
    canonical form. Unitaries are applied onto pairs of sites (i,i+1),
    (i+2,i+3), etc. We start at 'i=start' and move in increasing or
    decreasing order of sites depending on 'direction'

    Arguments:
    U         -- List of pairwise unitaries
    ψ         -- State in canonical form
    start     -- First site for applying pairwise unitaries
    direction -- Direction of sweep.

    Returns:
    ψ         -- MPS in canonical form"""

    #print("Apply pairwise unitarities in direction {} and at starting site {} with center {}".format(direction, start, ψ.center))
    ψ.recenter(start)
    if direction > 0:
        newstart = ψ.size-2
        for j in range(start, ψ.size-1, +2):
            #print('Updating sites ({}, {}), center={}, direction={}'.format(j, j+1, ψ.center, direction))
            AA = np.einsum('ijk,klm,nrjl -> inrm', ψ[j], ψ[j+1], U[j])
            ψ.update_2site(AA, j, +1, tolerance=tol)
            if j < newstart:
                ψ.update_canonical(ψ[j+1], +1, tolerance=tol)
            #print("New center= {}, new direction = {}".format(ψ.center, direction))
        return newstart, -1
    else:
        newstart = 0
        for j in range(start, -1, -2):
            #print('Updating sites ({}, {}), center={}, direction={}'.format(j, j+1, ψ.center, direction))
            AA = np.einsum('ijk,klm,nrjl -> inrm', ψ[j], ψ[j+1], U[j])
            ψ.update_2site(AA, j, -1, tolerance=tol)
            if j > 0:
                ψ.update_canonical(ψ[j], -1, tolerance=tol)
            #print("New center= {}, new direction = {}".format(ψ.center, direction))
        return newstart, +1
```

```python
# file: mps/evolution.py

class TEBD_evolution(object):
    """TEBD_evolution is a class that continuously updates a quantum state ψ
    evolving it with a Hamiltonian H over intervals of time dt."""

    def __init__(self, ψ, H, dt, timesteps=1, order=1, tol=DEFAULT_TOLERANCE):
        """Create a TEBD algorithm to evolve a quantum state ψ with a fixed
        Hamiltonian H.

        Arguments:
        ψ         -- Quantum state to be updated. The class keeps a copy.
        H         -- NNHamiltonian for the evolution
        dt        -- Size of each Trotter step
        timesteps -- How many Trotter steps in each call to evolve()
        order     -- Order of the Trotter approximation (1 or 2)
        tol       -- Tolerance in MPS truncation
        """
        self.H = H
        self.dt = float(dt)
        self.timesteps = timesteps
        self.order = order
        self.tolerance = tol
        self.Udt = pairwise_unitaries(H, dt)
        if order == 2:
            self.Udt2 = pairwise_unitaries(H, dt/2)
        if not isinstance(ψ, mps.state.CanonicalMPS):
            ψ = mps.state.CanonicalMPS(ψ, center=0)
        else:
            ψ = ψ.copy()
        self.ψ = ψ
        if ψ.center <= 1:
            self.start = 0
            self.direction = +1
        else:
            self.start = ψ.size-2
            self.direction = -1

    def evolve(self, timesteps=None):
        """Update the quantum state with `timesteps` repetitions of the
        Trotter algorithms."""

        #print("Apply TEBD for {} timesteps in the order {}".format(self.timesteps, self.order))

        if timesteps is None:
            timesteps = self.timesteps
        for i in range(self.timesteps):
            #print(i)
            if self.order == 1:
                #print("Sweep in direction {} and at starting site {}".format(self.direction, self.start))
                self.start, self.direction = apply_pairwise_unitaries(self.Udt, self.ψ, self.start, self.direction, tol=self.tolerance)
                #print("Sweep in direction {} and at starting site {}".format(self.direction, self.start))
                self.start, self.direction = apply_pairwise_unitaries(self.Udt, self.ψ, self.start, self.direction, tol=self.tolerance)
            else:
                self.start, self.direction = apply_pairwise_unitaries(self.Udt2, self.ψ, self.start, self.direction, tol=self.tolerance)
                self.start, self.direction = apply_pairwise_unitaries(self.Udt, self.ψ, self.start, self.direction, tol=self.tolerance)
                self.start, self.direction = apply_pairwise_unitaries(self.Udt2, self.ψ, self.start, self.direction, tol=self.tolerance)

            #print("New direction = {} and new starting site = {}".format(self.direction, self.start))

        return self.ψ

    def state():
        return self.ψ
```

## Error in Suzuki-Trotter decomposition

In the first order Suzuki-Trotter decomposition, evolution operator becomes
\begin{equation}
e^{-i \hat{H} \Delta t} = e^{-i \hat{H}_{\text{odd}} \Delta t}  e^{-i \hat{H}_{\text{even}} \Delta t} + O(\Delta t^2).
\end{equation}
Note that after $T/\Delta t$ time steps, the accumulated error is in the order of $\Delta t$.
Higher order Suzuki-Trotter decompositions can be used to reduce error.

# Tests

```python
# file: mps/test/test_TEBD.py
import unittest
import scipy.sparse as sp
import scipy.sparse.linalg
from mps.state import CanonicalMPS
from mps.tools import *
from mps.test.tools import *
from mps.evolution import *
from mps.hamiltonians import make_ti_Hamiltonian, ConstantNNHamiltonian

def random_wavefunction(n):
    ψ = np.random.rand(n) - 0.5
    return ψ / np.linalg.norm(ψ)

class TestTEBD_sweep(unittest.TestCase):

    @staticmethod
    def hopping_model(N, t, ω):
        a = annihilation(2)
        ad = creation(2)
        return make_ti_Hamiltonian(N, [t*a, t*ad], [ad, a], local_term = ω*(ad@a))

    @staticmethod
    def hopping_model_Trotter_matrix(N, t, ω):
        #
        # Hamiltonian that generates the evolution of the odd hoppings
        # and local frequencies
        return sp.diags([[t,0]*(N//2), [ω]+[ω/2]*(N-2)+[ω], [t,0]*(N//2)],
                        offsets=[-1,0,+1], shape=(N,N), dtype=np.float64)

    @staticmethod
    def hopping_model_matrix(N, t, ω):
        return sp.diags([[t]*(N), ω, [t]*(N)], offsets=[-1,0,+1], shape=(N,N))

    def inactive_test_apply_pairwise_unitaries(self):
        N = 2
        tt = -np.pi/2
        ω = np.pi
        dt = 0.1
        #
        # Numerically exact solution using Scipy's exponentiation routine
        ψwave = random_wavefunction(N)
        print(mps.state.wavepacket(ψwave).tovector())
        HMat = self.hopping_model_Trotter_matrix(N, tt, ω)
        ψwave_final = sp.linalg.expm_multiply(+1j * dt * HMat, ψwave)
        print(mps.state.wavepacket(ψwave_final).tovector())
        print(HMat.todense())
        #
        # Evolution using Trrotter
        H = self.hopping_model(N, tt, ω)
        U = pairwise_unitaries(H, dt)
        ψ = CanonicalMPS(mps.state.wavepacket(ψwave))
        start = 0
        direction = 1
        apply_pairwise_unitaries(U, ψ, start, direction, tol=DEFAULT_TOLERANCE)
        print(ψ.tovector())
        print(np.abs(mps.state.wavepacket(ψwave_final).tovector() - ψ.tovector()))

        self.assertTrue(similar(abs(mps.state.wavepacket(ψwave_final).tovector()),
                                abs(ψ.tovector())))

    def test_TEBD_evolution_first_order(self):
        #
        #
        #
        N = 19
        t = - np.pi/2
        ω = np.pi
        dt = 1e-6
        Nt = int(1000)
        #ψwave = random_wavefunction(N)
        xx=np.arange(N)
        x0 = int(N//2)
        w0 = 5
        k0 = np.pi/2
        #
        # Approximate evolution of a wavepacket in a tight-binding model
        ψwave = np.exp(-(xx-x0)**2 / w0**2 + 1j * k0*xx)
        ψwave = ψwave / np.linalg.norm(ψwave)
        Hmat = self.hopping_model_matrix(N, t, ω)
        ψwave_final = sp.linalg.expm_multiply(-1j * dt* Nt * Hmat, ψwave)
        #
        # Trotter solution
        ψmps = CanonicalMPS(mps.state.wavepacket(ψwave))
        H = self.hopping_model(N, t, ω)
        ψmps = TEBD_evolution(ψmps, H, dt, timesteps=Nt, order=1, tol=DEFAULT_TOLERANCE).evolve()

        self.assertTrue(similar(abs(mps.state.wavepacket(ψwave_final).tovector()),
                                abs(ψmps.tovector())))

    def test_TEBD_evolution_second_order(self):
        #
        #
        #
        N = 21
        t = 0.1
        ω = 0.5
        dt = 1e-6
        Nt = int(1000)
        #ψwave = random_wavefunction(N)
        xx=np.arange(N)
        x0 = int(N//2)
        w0 = 5
        k0 = np.pi/2
        #
        # Approximate evolution of a wavepacket in a tight-binding model
        ψwave = np.exp(-(xx-x0)**2 / w0**2 + 1j * k0*xx)
        ψwave = ψwave / np.linalg.norm(ψwave)
        Hmat = self.hopping_model_matrix(N, t, ω)
        ψwave_final = sp.linalg.expm_multiply(-1j * dt * Nt * Hmat, ψwave)
        #
        # Trotter evolution
        H = self.hopping_model(N, t, ω)
        ψmps = CanonicalMPS(mps.state.wavepacket(ψwave))
        ψmps = TEBD_evolution(ψmps, H, dt, timesteps=Nt, order=2, tol=DEFAULT_TOLERANCE).evolve()

        self.assertTrue(similar(abs(mps.state.wavepacket(ψwave_final).tovector()),
                                abs(ψmps.tovector())))
```

```python
suite1 = unittest.TestLoader().loadTestsFromNames(['__main__.TestTEBD_sweep'])
unittest.TextTestRunner(verbosity=2).run(suite1);
```


# Training tutorial

## Introduction to symbols

```python
from lambeq.ccg2discocat import DepCCGParser
from lambeq.tensor import TensorAnsatz
from lambeq.core.types import AtomicType
from discopy import Dim

# Define atomic types
N = AtomicType.NOUN
S = AtomicType.SENTENCE

# Parse a sentence
parser = DepCCGParser()
diagram = parser.sentence2diagram("John walks in the park")

# Apply a tensor ansatz
ansatz = TensorAnsatz({N: Dim(4), S: Dim(2)})
tensor_diagram = ansatz(diagram)
tensor_diagram.draw(figsize=(12,5), fontsize=12)
```

```python
tensor_diagram.free_symbols
```

    {John__n, in__s.r@n.r.r@n.r@s@n.l, park__n, the__n@n.l, walks__n.r@s}

```python
[s.size for s in tensor_diagram.free_symbols]
```

    [16, 256, 4, 8, 4]

```python
from lambeq.circuit import IQPAnsatz

iqp_ansatz = IQPAnsatz({N: 1, S: 1}, n_layers=1)
circuit = iqp_ansatz(diagram)
circuit.draw(figsize=(12,8), fontsize=12)
```

```python
[(s, s.size) for s in circuit.free_symbols]
```

    [(in__s.r@n.r.r@n.r@s@n.l_3, 1),
     (John__n_0, 1),
     (park__n_2, 1),
     (in__s.r@n.r.r@n.r@s@n.l_0, 1),
     (John__n_2, 1),
     (park__n_1, 1),
     (John__n_1, 1),
     (walks__n.r@s_0, 1),
     (the__n@n.l_0, 1),
     (in__s.r@n.r.r@n.r@s@n.l_1, 1),
     (in__s.r@n.r.r@n.r@s@n.l_2, 1),
     (park__n_0, 1)]

```python
from sympy import default_sort_key

parameters = sorted(tensor_diagram.free_symbols, key=default_sort_key)
```

```python
import numpy as np

tensors = [np.random.rand(p.size) for p in parameters]
print(tensors[0])
```

    [0.98111131 0.88772958 0.65068015 0.35379704]

```python
tensor_diagram_np = tensor_diagram.lambdify(*parameters)(*tensors)
print("Before lambdify:", tensor_diagram.boxes[0].data)
print("After lambdify:", tensor_diagram_np.boxes[0].data)
```

    Before lambdify: John__n
    After lambdify: [0.98111131 0.88772958 0.65068015 0.35379704]

```python
result = tensor_diagram_np.eval()
print(result)
```

    Tensor(dom=Dim(1), cod=Dim(2), array=[27.50087004, 24.96694656])

```python
result.array
```

    array([27.50087004, 24.96694656])

## A complete use case

```python
from discopy.tensor import Tensor
from jax import numpy as np
import numpy

np.random = numpy.random
Tensor.np = np
```

```python
# Read data
def read_data(fname):
    with open(fname, 'r') as f:
        lines = f.readlines()
    data, targets = [], []
    for ln in lines:
        t = int(ln[0])
        data.append(ln[1:].strip())
        targets.append(np.array([t, not(t)], dtype=np.float32))
    return data, np.array(targets)

train_data, train_targets = read_data('examples/datasets/mc_train_data.txt')
test_data, test_targets = read_data('examples/datasets/mc_test_data.txt')
dev_data, dev_targets = read_data('examples/datasets/mc_dev_data.txt')
```

    WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

```python
train_data[:10]
```

    ['skillful man prepares sauce',
     'skillful man bakes dinner',
     'woman cooks tasty meal',
     'man prepares meal',
     'skillful woman debugs program',
     'woman prepares tasty meal',
     'person runs program',
     'person runs useful application',
     'woman prepares sauce',
     'woman prepares dinner']

```python
train_targets
```

    DeviceArray([[1., 0.],
                 [1., 0.],
                 [1., 0.],
                 ...,
                 [0., 1.],
                 [1., 0.],
                 [0., 1.]], dtype=float32)

```python
# Parse sentences to diagrams

from lambeq.ccg2discocat import DepCCGParser

parser = DepCCGParser()
train_diagrams = parser.sentences2diagrams(train_data)
test_diagrams = parser.sentences2diagrams(test_data)

train_diagrams[0].draw(figsize=(8,4), fontsize=13)
```

```python
# Create ansatz and convert to tensor diagrams

from lambeq.tensor import SpiderAnsatz
from lambeq.core.types import AtomicType
from discopy import Dim

N = AtomicType.NOUN
S = AtomicType.SENTENCE

# Create an ansatz by assigning 2 dimensions to both
# noun and sentence spaces
ansatz = SpiderAnsatz({N: Dim(2), S: Dim(2)})

train_circuits = [ansatz(d) for d in train_diagrams]
test_circuits = [ansatz(d) for d in test_diagrams]

all_circuits = train_circuits + test_circuits

all_circuits[0].draw(figsize=(8,4), fontsize=13)
```

```python
# Create vocabulary

from sympy import default_sort_key

vocab = sorted(
   {sym for circ in all_circuits for sym in circ.free_symbols},
    key=default_sort_key
)
tensors = [np.random.rand(w.size) for w in vocab]

tensors[0]
```

    array([0.35743395, 0.45764418])

```python
def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def loss(tensors):
    # Lambdify
    np_circuits = [c.lambdify(*vocab)(*tensors) for c in train_circuits]
    # Compute predictions
    predictions =  sigmoid(np.array([c.eval().array for c in np_circuits]))

    # binary cross-entropy loss
    cost = -np.sum(train_targets * np.log2(predictions)) / len(train_targets)
    return cost
```

```python
from jax import jit, grad

training_loss = jit(loss)
gradient = jit(grad(loss))
```

```python
training_losses = []

epochs = 90

for i in range(epochs):

    gr = gradient(tensors)
    for k in range(len(tensors)):
        tensors[k] = tensors[k] - gr[k] * 1.0

    training_losses.append(float(training_loss(tensors)))

    if (i + 1) % 10 == 0:
        print(f"Epoch {i + 1} - loss {training_losses[-1]}")
```

    Epoch 10 - loss 0.07233709841966629
    Epoch 20 - loss 0.015333528630435467
    Epoch 30 - loss 0.00786149874329567
    Epoch 40 - loss 0.00515687046572566
    Epoch 50 - loss 0.0037753921933472157
    Epoch 60 - loss 0.0029438300989568233
    Epoch 70 - loss 0.002392344642430544
    Epoch 80 - loss 0.0020021884702146053
    Epoch 90 - loss 0.001713048666715622

```python
# Testing

np_test_circuits = [c.lambdify(*vocab)(*tensors) for c in test_circuits]
test_predictions =  sigmoid(np.array([c.eval().array for c in np_test_circuits]))

hits = 0
for i in range(len(np_test_circuits)):
    target = test_targets[i]
    pred = test_predictions[i]
    if np.argmax(target) == np.argmax(pred):
        hits += 1

print("Accuracy on test set:", hits / len(np_test_circuits))
```

    Accuracy on test set: 0.9


# The Tractable Buffer Stock Model

```python
# This cell has just a bit of initial setup. You can click the arrow to expand it.
%matplotlib inline
import matplotlib.pyplot as plt

import sys
import os
sys.path.insert(0, os.path.abspath('../lib'))

import numpy as np
import HARK
from time import clock
from copy import deepcopy
mystr = lambda number : "{:.3f}".format(number)

from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets

from HARK.utilities import plotFuncs
from HARK.ConsumptionSaving.TractableBufferStockModel import TractableConsumerType
```

The [TractableBufferStock](http://www.econ2.jhu.edu/people/ccarroll/public/LectureNotes/Consumption/TractableBufferStock/) model is a (relatively) simple framework that captures all of the qualitative, and many of the quantitative features of optimal consumption in the presence of labor income uncertainty.

The key assumption behind the model's tractability is that there is only a single, stark form of uncertainty.  So long as an employed consumer remains employed, his labor income will rise at a constant rate.  But between any period and the next there is constant hazard $p$ of transitioning to the "unemployed" state. Unemployment is irreversible, like retirement or disability.  When unemployed, the consumer receives a fixed amount of income (for simplicity, zero).

\begin{eqnarray*}
\mathbb{E}_{t}[V_{t+1}^{\bullet}(M_{t+1})] & = & (1-p)V_{t+1}^{e}(M_{t+1})+p V_{t+1}^{u}(M_{t+1})
\end{eqnarray*}

A consumer with CRRA utility $U(C) = \frac{C^{1-\rho}}{1-\rho}$ solves an optimization problem that looks standard (where $P$ is Permanent income and $\Gamma$ is the income growth factor):

\begin{eqnarray*}
V_t(M_t) &=& \max_{C_t} ~ U(C_t) + \beta \mathbb{E}[V_{t+1}^{\bullet}], \\
M_{t+1} &=& R A_t + \mathbb{1}(P_{t+1}), \\
P_{t+1} &=& \Gamma_{t+1} P_t,
\end{eqnarray*}

where $\mathbb{1}$ is an indicator of whether the consumer is employed in the next period.

Under plausible parameter values the model has a target level of $m = M/P$ (market resources to permanent income) with an analytical solution that exhibits plausible relationships among all of the parameters.  (See the linked handout for details).

```python
# Define a parameter dictionary and representation of the agents for the tractable buffer stock model
TBS_dictionary =  {'UnempPrb' : .00625,    # Probability of becoming unemployed
                   'DiscFac' : 0.975,      # Intertemporal discount factor
                   'Rfree' : 1.01,         # Risk-free interest factor on assets
                   'PermGroFac' : 1.0025,  # Permanent income growth factor (uncompensated)
                   'CRRA' : 1.5}           # Coefficient of relative risk aversion
MyTBStype = TractableConsumerType(**TBS_dictionary)
```

## Target Wealth

Whether the model exhibits a "target" or "stable" level of wealth for employed consumers depends on whether the 'Growth Impatience Condition' (the GIC) holds:

\begin{equation}\label{eq:GIC}
 \left(\frac{(R \beta (1-\mho))^{1/\rho}}{\Gamma}\right)  <  1
\end{equation}

### IMPORTENT: To activate: jupyter-js-widgets/extension in the nbextensions

```python
# Make an interactive plot of the tractable buffer stock solution

# To make some of the widgets not appear, replace X_widget with fixed(desired_fixed_value) in the arguments below.
interact(makeTBSplot,
         DiscFac = DiscFac_widget,
#         CRRA = CRRA_widget,
         CRRA = fixed(2.5),
         Rfree = Rfree_widget,
         PermGroFac = PermGroFac_widget,
         UnempPrb = UnempPrb_widget,
         mMin = mMin_widget,
         mMax = mMax_widget,
         cMin = cMin_widget,
         cMax = cMax_widget,
         show_targ = show_targ_widget,
         plot_emp = plot_emp_widget,
         plot_ret = plot_ret_widget,
         plot_mSS = plot_mSS_widget,
        );

```


# Two coupled equilibria: protolysis of ammonia in water
In this notebook we will look at how ``ChemPy`` can be used to formulate a system of (non-linear) equations from conservation laws and equilibrium equations. We will look att ammonia since it is a fairly well-known subtance.

```python
from operator import mul
from functools import reduce
from itertools import product
import chempy
from chempy.chemistry import Species, Equilibrium
from chempy.equilibria import EqSystem, NumSysLog, NumSysLin
import numpy as np
import sympy as sp
sp.init_printing()
import matplotlib.pyplot as plt
%matplotlib inline
sp.__version__, chempy.__version__
```

First, let's define our substances. ChemPy can parse chemical formulae into ``chempy.Substance`` instances with information on charge, composition, molar mass (and pretty-printing):

```python
substance_names = ['H+', 'OH-', 'NH4+', 'NH3', 'H2O']
subst = {n: Species.from_formula(n) for n in substance_names}
assert [subst[n].charge for n in substance_names] == [1, -1, 1, 0, 0], "Charges of substances"
print(u'Composition of %s: %s' % (subst['NH3'].unicode_name, subst['NH3'].composition))
```

Let's define som e initial concentrations and governing equilibria:

```python
init_conc = {'H+': 1e-7, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}
x0 = [init_conc[k] for k in substance_names]
H2O_c = init_conc['H2O']
w_autop = Equilibrium({'H2O': 1}, {'H+': 1, 'OH-': 1}, 10**-14/H2O_c)
NH4p_pr = Equilibrium({'NH4+': 1}, {'H+': 1, 'NH3': 1}, 10**-9.26)
equilibria = w_autop, NH4p_pr
[(k, init_conc[k]) for k in substance_names]
```

```python
eqsys = EqSystem(equilibria, subst)
eqsys
```

We can solve the non-linear system of equations by calling the ``root`` method (the underlying representation uses [pyneqsys](https://github.com/bjodah/pyneqsys):

```python
x, sol, sane = eqsys.root(init_conc)
x, sol['success'], sane
```

This system is quite easy to solve for, if we have convergence problems we can try to solve a transformed system. As an example we will use ``NumSysLog``:

```python
logx, logsol, sane = eqsys.root(init_conc, NumSys=(NumSysLog,))
logx, logsol['success'], sane
```

In thise case they give essentially the same answer:

```python
x - logx
```

We can create symbolic representations of these systems of equations:

```python
ny = len(substance_names)
y = sp.symarray('y', ny)
i = sp.symarray('i', ny)
K = Kw, Ka = sp.symbols('K_w K_a')
w_autop.param = Kw
NH4p_pr.param = Ka
ss = sp.symarray('s', ny)
ms = sp.symarray('m', ny)
```

```python
numsys_log = NumSysLog(eqsys, backend=sp)
f = numsys_log.f(y, list(i)+list(K))
f
```

```python
numsys_lin = NumSysLin(eqsys, backend=sp)
numsys_lin.f(y, i)
```

```python
A, ks = eqsys.stoichs_constants(False, backend=sp)
[reduce(mul, [b**e for b, e in zip(y, row)]) for row in A]
```

```python
from pyneqsys.symbolic import SymbolicSys
subs = list(zip(i, x0)) + [(Kw, 10**-14), (Ka, 10**-9.26)]
numf = [_.subs(subs) for _ in f]
neqs = SymbolicSys(list(y), numf)
neqs.solve([0, 0, 0, 0, 0], solver='scipy')
```

```python
j = sp.Matrix(1, len(f), lambda _, q: f[q]).jacobian(y)
init_conc_j = {'H+': 1e-10, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}
xj = eqsys.as_per_substance_array(init_conc_j)
jarr = np.array(j.subs(dict(zip(y, xj))).subs({Kw: 1e-14, Ka: 10**-9.26}).subs(
            dict(zip(i, xj))))
jarr = np.asarray(jarr, dtype=np.float64)
np.log10(np.linalg.cond(jarr))
```

```python
j.simplify()
j
```

```python
eqsys.composition_balance_vectors()
```

```python
numsys_rref_log = NumSysLog(eqsys, True, True, backend=sp)
numsys_rref_log.f(y, list(i)+list(K))
```

```python
np.set_printoptions(4, linewidth=120)
scaling = 1e8
for rxn in eqsys.rxns:
    rxn.param = rxn.param.subs({Kw: 1e-14, Ka: 10**-9.26})
```

```python
x, res, sane = eqsys.root(init_conc, rref_equil=True, rref_preserv=True)
x, res['success'], sane
```

```python
x, res, sane = eqsys.root(init_conc, x0=eqsys.as_per_substance_array(
        {'H+': 1e-11, 'OH-': 1e-3, 'NH4+': 1e-3, 'NH3': 1.0, 'H2O': 55.5}))
res['success'], sane
```

```python
x, res, sane = eqsys.root(init_conc, x0=eqsys.as_per_substance_array(
        {'H+': 1.7e-11, 'OH-': 3e-2, 'NH4+': 3e-2, 'NH3': 0.97, 'H2O': 55.5}))
res['success'], sane
```

```python
init_conc
```

```python
nc=60
Hp_0 = np.logspace(-3, 0, nc)
def plot_rref(**kwargs):
    fig, axes = plt.subplots(2, 2, figsize=(16, 6), subplot_kw=dict(xscale='log', yscale='log'))
    return [eqsys.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': axes.flat[i]}, rref_equil=e,
                        rref_preserv=p, **kwargs) for i, (e, p) in enumerate(product(*[[False, True]]*2))]
```

```python
res_lin = plot_rref(method='lm')
[all(_[2]) for _ in res_lin]
```

```python
for col_id in range(len(substance_names)):
    for i in range(1, 4):
        plt.subplot(1, 3, i, xscale='log')
        plt.gca().set_yscale('symlog', linthreshy=1e-14)
        plt.plot(Hp_0, res_lin[0][0][:, col_id] - res_lin[i][0][:, col_id])
plt.tight_layout()
```

```python
eqsys.plot_errors(res_lin[0][0], init_conc, Hp_0, 'H+')
```

```python
init_conc, eqsys.ns
```

```python
res_log = plot_rref(NumSys=NumSysLog)
```

```python
eqsys.plot_errors(res_log[0][0], init_conc, Hp_0, 'H+')
```

```python
res_log_lin = plot_rref(NumSys=(NumSysLog, NumSysLin))
eqsys.plot_errors(res_log_lin[0][0], init_conc, Hp_0, 'H+')
```

```python
from chempy.equilibria import NumSysSquare
res_log_sq = plot_rref(NumSys=(NumSysLog, NumSysSquare))
eqsys.plot_errors(res_log_sq[0][0], init_conc, Hp_0, 'H+')
```

```python
res_sq = plot_rref(NumSys=(NumSysSquare,), method='lm')
eqsys.plot_errors(res_sq[0][0], init_conc, Hp_0, 'H+')
```

```python
x, res, sane = eqsys.root(x0, NumSys=NumSysLog, rref_equil=True, rref_preserv=True)
x, res['success'], sane
```

```python
x, res, sane = eqsys.root(x, NumSys=NumSysLin, rref_equil=True, rref_preserv=True)
x, res['success'], sane
```


# Physics 420/580 Final Exam
## December 19, 2017 2pm-5pm

Do three of the four following problems.  If you attempt all four, your best three will be used for your mark.  Use the Jupyter notebook, inserting your code and any textual answers/explanations in cells between the questions. (Feel free to add additional cells!)  Marks will be given based on how clearly you demonstrate your understanding.

There are no restrictions on downloading from the internet, eclass,  or the use of books, notes, or any other widely available computing resources.  However, **you are not allowed** to communicate with each other or collaborate in any way and uploading to the internet or sending or receiving direct communications is not appropriate.

When you are finished, upload the jupyter notebook to eclass. Eclass times out after 5:05 so make sure that you upload things before then.  Also be careful to save the notebook periodically and that you upload your exam file.

```python

```

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

# Monte Carlo

Write a Monte Carlo that calculates that simulates a game in which a person throws a pair of dice 1000 times.  Histogram the outcomes (ie. the sum of the two dice) for one game.  Define the "streak length" by the largest number of consecutive throws where the number thrown doesn't decrease.  For example, if somewhere in the sequence you threw 2,4,5,7,7,11  (one after the other) that would be a streak of length 6. Repeat the game 500 times (rolling the dice 1000 times inside each game), and plot the distribution of maximum streak lengths.

```python
dice_sum = []
for i in range(1000):
    dice_sum.append(np.random.randint(1, 7) + np.random.randint(1, 7))
```

```python
plt.hist(np.array(dice_sum), bins=50)
```

```python
class Game:
    def __init__(self):
        self.streak = 0
        self.max_number = 0
        self.max_streak = 0

    def throw_dice(self, throw):
        if self.streak >= self.max_streak:
            self.max_streak = self.streak
        if throw >= self.max_number:
            self.streak += 1
            self.max_number = throw
        else:
            self.streak = 0
            self.max_number = 0
```

```python
streak_lengths = []
for j in range(500):
    game = Game()
    for i in range(1000):
        dice_sum = (np.random.randint(1, 7) + np.random.randint(1, 7))
        game.throw_dice(dice_sum)
    streak_lengths.append(game.max_streak)
```

```python
plt.hist(np.array(streak_lengths),bins=50)
```

# Fitting

```python

```

Download any one of the datasets in the "Final Exam Data Sets" folder on eclass- they have names like DataSet####.npy.  The data sets are npy arrays, which can be read in with the np.load command and consists of one dimensional data (x-values) that have been thrown according to a distribution.

Histogram and plot the data using 100 bins.  Make the title of the histogram include the name of the dataset.
Make an array y that contains the histogram values (counts per bin).
Make a corresponding array of the central x-value for each bin.is

```python
x = np.load('Dataset32.npy')
bins = plt.hist(x, 100)
counts = bins[0]
edges = bins[1]
```

```python
counts.shape
```

    (100,)

```python
centers = []
for i in range(99):
    print(i, counts[i], (counts[i+1] + counts[i])/2)
    centers.append((counts[i+1] + counts[i])/2)
```

    0 672314.0 647652.5
    1 622991.0 598915.0
    2 574839.0 553323.0
    3 531807.0 511180.0
    4 490553.0 471737.5
    5 452922.0 434471.5
    6 416021.0 399040.0
    7 382059.0 366627.5
    8 351196.0 337261.5
    9 323327.0 309211.0
    10 295095.0 282284.0
    11 269473.0 258361.5
    12 247250.0 236282.0
    13 225314.0 215906.0
    14 206498.0 197766.5
    15 189035.0 180123.0
    16 171211.0 163404.5
    17 155598.0 148299.5
    18 141001.0 135150.0
    19 129299.0 123489.0
    20 117679.0 112691.5
    21 107704.0 102859.5
    22 98015.0 93878.0
    23 89741.0 85870.0
    24 81999.0 78723.5
    25 75448.0 72168.5
    26 68889.0 66735.0
    27 64581.0 61985.5
    28 59390.0 57350.5
    29 55311.0 53580.5
    30 51850.0 50143.0
    31 48436.0 47373.5
    32 46311.0 45188.5
    33 44066.0 43270.5
    34 42475.0 41726.5
    35 40978.0 40325.5
    36 39673.0 38933.5
    37 38194.0 38070.5
    38 37947.0 37645.5
    39 37344.0 37150.0
    40 36956.0 36739.5
    41 36523.0 36504.0
    42 36485.0 36696.5
    43 36908.0 36689.0
    44 36470.0 36823.5
    45 37177.0 37527.0
    46 37877.0 37940.0
    47 38003.0 38182.0
    48 38361.0 38782.5
    49 39204.0 39577.0
    50 39950.0 40208.5
    51 40467.0 40952.5
    52 41438.0 42100.5
    53 42763.0 42988.0
    54 43213.0 43772.5
    55 44332.0 45079.0
    56 45826.0 46716.5
    57 47607.0 48553.5
    58 49500.0 50383.5
    59 51267.0 52165.0
    60 53063.0 54533.5
    61 56004.0 57413.5
    62 58823.0 60610.0
    63 62397.0 64048.0
    64 65699.0 67764.5
    65 69830.0 71869.0
    66 73908.0 76642.0
    67 79376.0 82266.0
    68 85156.0 88396.0
    69 91636.0 95238.0
    70 98840.0 102699.5
    71 106559.0 110615.5
    72 114672.0 119536.5
    73 124401.0 129657.5
    74 134914.0 141161.0
    75 147408.0 153414.5
    76 159421.0 166432.0
    77 173443.0 181218.0
    78 188993.0 196790.0
    79 204587.0 213834.0
    80 223081.0 232554.0
    81 242027.0 252623.0
    82 263219.0 274322.5
    83 285426.0 297834.0
    84 310242.0 323525.5
    85 336809.0 351381.5
    86 365954.0 381086.5
    87 396219.0 412915.5
    88 429612.0 446650.0
    89 463688.0 482533.5
    90 501379.0 520856.5
    91 540334.0 561596.5
    92 582859.0 606432.5
    93 630006.0 653781.0
    94 677556.0 702979.0
    95 728402.0 755911.5
    96 783421.0 812132.0
    97 840843.0 871345.5
    98 901848.0 934319.0

Write a method using numpy.poly1d to describe a fourth order (five parameter) polynomials.  Use this method to plot \begin{equation} 4x^4+3x^3+2x^2+x+10\end{equation} between -2 and 2.

```python
from scipy.interpolate import interp1d
```

```python
xdata = np.linspace(-2, 2, 99)
func = interp1d(xdata, centers)
data = np.linspace(-2, 2)
plt.plot(func(data))
```

Write a method poly5residuals(parameters,*args), where "parameters" is an numpy array that contains the polynomial coefficients, and args will contain three np arrays- x,y,sigma that contain the x values, y values, and uncertainties of the data.

poly5residuals should return an array of residuals: each element corresponds to a data element and is (polynomial(x)-y)/sigma.  It should ignore any zero values- either delete the corresponding bin values or make the residuals for zero values in the histogram should be zero.

Print out a list of residuals between the polynomial and the histogram (the x values should correspond to the center of each bin).  Sigma should correspond to the uncertainties in the data for a counting experiment of uncorrelated events.

```python
poly = np.poly1d([4, 3, 2, 1, 10])
plt.plot(xdata, poly(xdata))
```

```python
from scipy.optimize import least_squares
```

```python
def poly5residuals(parameters, *args):
    xdata, ydata, sigma = args
    poly = np.poly1d(parameters)

    return (poly(xdata) - ydata)/sigma
```

```python
poly5residuals(startingParam, *args)
```

    array([-804.69901506, -773.82867841, -743.79172655, -714.90555639,
           -686.77056728, -659.08552318, -631.63891598, -605.44224559,
           -580.6884059 , -556.01517594, -531.25312729, -508.24312242,
           -486.04007307, -464.60956321, -444.66320882, -424.36356849,
           -404.18840794, -385.05243438, -367.5837013 , -351.36659513,
           -335.65236654, -320.67408954, -306.35204999, -292.99256621,
           -280.53356189, -268.59801964, -258.28686545, -248.92411606,
           -239.43420349, -231.42856472, -223.87954878, -217.60712556,
           -212.5279139 , -207.96710922, -204.22164535, -200.76256061,
           -197.26555641, -195.06605841, -193.9737409 , -192.69234156,
           -191.62423984, -191.00876201, -191.51196872, -191.49232854,
           -191.84311645, -193.66786759, -194.73108823, -195.35132992,
           -196.88214219, -198.88942286, -200.47048107, -202.31734713,
           -205.13455283, -207.28612469, -209.16945952, -212.26906466,
           -216.0912103 , -220.30007482, -224.41431497, -228.34817625,
           -233.47575875, -239.56293057, -246.14281297, -253.02881424,
           -260.26770342, -268.03519175, -276.79383996, -286.7710998 ,
           -297.26478729, -308.55613394, -320.41653035, -332.53711942,
           -345.68760823, -360.02596929, -375.6590862 , -391.62553345,
           -407.90331054, -425.63821415, -443.54974598, -462.35976961,
           -482.17466684, -502.55041224, -523.69032818, -545.67189639,
           -568.72108507, -592.70026259, -617.24556147, -642.50610608,
           -668.23732195, -694.5621833 , -721.61700154, -749.30768431,
           -778.64518776, -808.47173101, -838.33979649, -869.32992924,
           -901.07865785, -933.35055173, -966.49004779])

Fit the polynomial to the histogram using scipy.optimize.least_squares.
Plot the histogram with the polynomial overlaid and separately plot the residuals.

```python
xdata = np.linspace(-2, 2, 99)
ydata = centers
sigma = np.sqrt(ydata)

args = (xdata, ydata, sigma )

out = least_squares(poly5residuals, startingParam, args=args)
```

Extract the marginalized uncertainties for each parameter from the fit.

```python
out.x
plt.plot(xdata, np.poly1d(out.x)(xdata))

```

## Differential Equations
The simple electronic circuit below can be described by the three differential equations:
\begin{equation}V_i-i_1R_1-\frac{q_1}{C}=0\end{equation}
\begin{equation}-i_2R_2=-\frac{q_3}{C}=V_o\end{equation}
\begin{equation}i_1-i_2-i_3=0\end{equation}

Where $i_{j}=\frac{dq_{j}}{dt}$ for j=1,2,3; $V_i$ is the input voltage and $V_o$ is the output voltage.Use $R_1=R_2=100000$ ohms and C=25e-9 farad.

The diagram is just for reference/interest.  You only need the differential equations to do the problems.

Write a program that solves for the currents, $i_{1,2,3}(t)$ given $q_{1,2,3}(t)$ and a function $V_i(t)$.

```python
C = 25e-9
R1 = 100000
R2 = 100000
frequency = 1000
omega = 2*np.pi*frequency
```

```python
def solve_currents(q, Vi):
    dydt = np.empty(3)

    dydt[0] = (1/R1)*(Vi-q[0]/C)
    dydt[1] = (1/R2)*(q[2]/C)
    dydt[2] = (1/R1)*(Vi-q[0]/C) - (1/R2)*(q[2]/C)

    return dydt
```

 Suppose $V_i=5\sin\omega t$.  If the frequency $\nu=1000$ Hz, the angular frequency is $\omega=2\pi\nu=2\pi 1000$.  Using odeint, solve the differential equation and plot $V_o(t)$, over 20 periods of oscillation.  Repeat for 100 Hz and 5000 Hz.

```python
def dydt(y, t):
    Vi = 5*np.sin(omega*t)
    q = y
    return solve_currents(q, Vi)
```

```python
from scipy.integrate import odeint
t = np.linspace(0, 20/frequency, 1000)
```

```python
q0 = np.array([0, 0, 0])

y = odeint(dydt, q0, t,rtol=1e-13,atol=1e-13)
```

```python
plt.plot(t, y[:, 1]*R2*(-1))
```

Plot the amplitude of $V_o$ as a function of $\omega$ for frequencies from 10 Hz to 10 kHz.

```python
for frequency in [100, 1000, 5000]:
    omega = 2*np.pi*frequency

    y = odeint(dydt, q0, t,rtol=1e-13,atol=1e-13)
    plt.plot(t, -y[:, 2]/C, label=f'{frequency}')
```

## Geometry and Monte Carlo
Using Monte Carlo, calculate the volume of an object made by molding 3 10 cm diameter cylinders together, so that the angle between any two axes is 60 degrees, and trimming the cylinders at +/- 30 cm with respect to each cylinder axis.

```python
def spherical_to_xyz(r, theta, phi):
    return np.array([r*np.sin(theta)*np.cos(phi), r*np.sin(theta)*np.sin(phi), r*np.cos(theta)])
```

```python
k1 = spherical_to_xyz(1, 0, 0)
k2 = spherical_to_xyz(1, np.pi/3, 0)
k3 = spherical_to_xyz(1, np.pi/3, 1.2309594)

print(np.dot(k1, k2), np.dot(k1, k3), np.dot(k2, k3))
```

    0.5000000000000001 0.5000000000000001 0.5000000122617794

When you are finished, upload your completed exam onto eClass.

```python
from scipy.optimize import minimize

def find_angle(x):
    k3 = spherical_to_xyz(1, np.pi/3, x)
    return (np.dot(k2, k3) - 0.5)**2

minimize(find_angle, 2)
```

          fun: 2.0222669039148207e-16
     hess_inv: array([[0.99876259]])
          jac: array([-1.26604457e-08])
      message: 'Optimization terminated successfully.'
         nfev: 12
          nit: 3
         njev: 4
       status: 0
      success: True
            x: array([1.2309594])

```python
throws = np.random.uniform(-50, 50, size=(100000, 3))

n_inside = 0
n_outside = 0
for i in range(100000):
    throw = throws[i]

    for k in [k1, k2, k3]:
        distance_normal_axis = throw - np.dot(throw, k)*k
        distance_along_axis = np.sqrt(np.dot(throw, k))
        if distance_along_axis < 30 and (np.linalg.norm(distance_normal_axis) < 5):
            n_inside += 1
            break
    else:
        n_outside += 1

print(n_inside, n_outside)
```

    C:\Users\Zhi\Anaconda3\envs\phys420\lib\site-packages\ipykernel_launcher.py:10: RuntimeWarning: invalid value encountered in sqrt
      # Remove the CWD from sys.path while we load stuff.

    1213 98787

```python
print(100**3*(n_inside/100000))
```

    12130.0

```python

```