# Five-Link biped walking loop problem: interactive demonstration

Hello and welcome. This is a Jupyter Notebook, a kind of document that can alternate between static content, like text and images, and executable cells of code.

This document ilustrates the Five-link biped walking loop test case of the paper: "Collocation Methods for Second Order Systems", submitted to RSS 2022.

In order to run the cells of code, you can select the cell and clic on the small "play" button in the bar above or press shift+enter. Alternatively, you can select the option "run" -> "run all cells" in order to run all the code in order. Beware that some cells can take several minutes!

All of the code used in this example is open-source and free to use.

[SymPy](https://www.sympy.org/en/index.html) is used for Symbolic formulation and manipulation of the problem.

[Numpy](https://numpy.org/) is used for numerical arrays and operations.

[CasADI](https://web.casadi.org/) is used for optimization.

[Optibot](https://github.com/AunSiro/optibot) is the name of the package where we are compiling our code. We aim to produce a toolbox for Optimal Control Problems, focused on robotics, including a high level, readable and clean interface between the prior three packages.

## Package imports

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

```python
from sympy import (symbols, simplify)
from sympy.physics.mechanics import dynamicsymbols, init_vprinting
from sympy.physics.mechanics import Lagrangian, ReferenceFrame, Point, Particle,inertia, RigidBody, angular_momentum
```

```python
from optibot.symbolic import lagrange, diff_to_symb, SimpLagrangesMethod
from optibot.numpy import unpack
```

```python
#SymPy vector-like latex rendering inizialization:

init_vprinting()
```

## Symbolic Problem Modelling

The first step is to model our problem taking advantage of the high level object syntax of the mechanics module in SymPy

```python
# Creating symbols and dynamic symbols

m0, m1, m2, m3, m4, l0, l1, l2, l3, l4, t, g = symbols('m_0:5 l_0:5 t g')
I0, I1, I2, I3, I4, d0, d1, d2, d3, d4 = symbols('I_0:5 d_0:5')
q0, q1, q2, q3, q4 = dynamicsymbols('q_0:5')
m0, m1, m2, m3, m4, l0, l1, l2, l3, l4, t, g, I0, I1, I2, I3, I4, d0, d1, d2, d3, d4, q0, q1, q2, q3, q4
```

```python
# Definition of the physics system

N_in = ReferenceFrame('N')
P0 = Point('P0')
P0.set_vel(N_in, 0)

N0 = N_in.orientnew('N0', 'Axis', [q0, N_in.z])
P1 = P0.locatenew('P1', l0 * N0.y)
P1.set_vel(N_in, P1.pos_from(P0).dt(N_in))
CM0 = P0.locatenew('CM0', (l0-d0) * N0.y)
CM0.set_vel(N_in, CM0.pos_from(P0).dt(N_in))
I_0 = inertia(N0, 0, 0, I0)
body0 = RigidBody('Stance_Tibia', CM0, N0, m0, (I_0,CM0))
body0.potential_energy = m0 * g * CM0.pos_from(P0).dot(N_in.y)

N1 = N_in.orientnew('N1', 'Axis', [q1, N_in.z])
P2 = P1.locatenew('P2', l1 * N1.y)
P2.set_vel(N_in, P2.pos_from(P0).dt(N_in))
CM1 = P1.locatenew('CM1', (l1-d1) * N1.y)
CM1.set_vel(N_in, CM1.pos_from(P0).dt(N_in))
I_1 = inertia(N1, 0, 0, I1)
body1 = RigidBody('Stance_Femur', CM1, N1, m1, (I_1,CM1))
body1.potential_energy = m1 * g * CM1.pos_from(P0).dot(N_in.y)

N2 = N_in.orientnew('N2', 'Axis', [q2, N_in.z])
P3 = P2.locatenew('P3', l2 * N2.y)
P3.set_vel(N_in, P3.pos_from(P0).dt(N_in))
CM2 = P2.locatenew('CM2', d2 * N2.y)
CM2.set_vel(N_in, CM2.pos_from(P0).dt(N_in))
I_2 = inertia(N2, 0, 0, I2)
body2 = RigidBody('Torso', CM2, N2, m2, (I_2,CM2))
body2.potential_energy = m2 * g * CM2.pos_from(P0).dot(N_in.y)

N3 = N_in.orientnew('N3', 'Axis', [q3, N_in.z])
P4 = P2.locatenew('P4', -l3 * N3.y)
P4.set_vel(N_in, P4.pos_from(P0).dt(N_in))
CM3 = P2.locatenew('CM3', -d3 * N3.y)
CM3.set_vel(N_in, CM3.pos_from(P0).dt(N_in))
I_3 = inertia(N3, 0, 0, I3)
body3 = RigidBody('Swing_Femur', CM3, N3, m3, (I_3,CM3))
body3.potential_energy = m3 * g * CM3.pos_from(P0).dot(N_in.y)

N4 = N_in.orientnew('N4', 'Axis', [q4, N_in.z])
P5 = P4.locatenew('P5', -l4 * N4.y)
P5.set_vel(N_in, P5.pos_from(P0).dt(N_in))
CM4 = P4.locatenew('CM4', -d4 * N4.y)
CM4.set_vel(N_in, CM4.pos_from(P0).dt(N_in))
I_4 = inertia(N4, 0, 0, I4)
body4 = RigidBody('Swing_Tibia', CM4, N4, m4, (I_4,CM4))
body4.potential_energy = m4 * g * CM4.pos_from(P0).dot(N_in.y)
```

```python
#Computing the Lagrangian

Lag_simp = Lagrangian(N_in, body0, body1, body2, body3, body4)
Lag_simp
```

```python
from optibot.symbolic import ImplicitLagrangesMethod
```

```python
# Defining the control forces and external actions, and applying them to our system

u0, u1, u2, u3, u4 = symbols('u_:5')
FL = [
    (N0, (u0-u1) * N_in.z),
    (N1, (u1-u2) * N_in.z),
    (N2, (u2-u3) * N_in.z),
    (N3, (u3-u4) * N_in.z),
    (N4, u4 * N_in.z)
]
LM_small = ImplicitLagrangesMethod(Lag_simp, [q0, q1, q2, q3, q4], forcelist=FL, frame=N_in)
```

```python
# Generating the dynamic equations

LM_small.form_lagranges_equations()
```

```python
impl_x = LM_small.implicit_dynamics_x
```

```python
impl_q = LM_small.implicit_dynamics_q
```

### Generating auxiliar functions

Later in the problem we need some expressions derived from the problem. Here we will generate them as symbolic expressions, and then convert them to numerical functions.

```python
import casadi as cas
from sympy import lambdify
from optibot.casadi import implicit_dynamic_q_to_casadi_function, implicit_dynamic_x_to_casadi_function, sympy2casadi
from optibot.symbolic import find_arguments, diff_to_symb_expr
from sympy.physics.mechanics import kinetic_energy, potential_energy
```

```python
imp_dyn_x_f_cas = implicit_dynamic_x_to_casadi_function(impl_x, list(dynamicsymbols('x_0:10')), verbose=True)
```

```python
imp_dyn_q_f_cas = implicit_dynamic_q_to_casadi_function(impl_q, list(LM_small.q), verbose=True)
```

```python
imp_dyn_x_f_cas
```

```python
feet_x = P5.pos_from(P0).dot(N_in.x)
feet_x = diff_to_symb_expr(feet_x)
feet_x
```

```python
feet_y = P5.pos_from(P0).dot(N_in.y)
feet_y = diff_to_symb_expr(feet_y)
feet_y
```

```python
feet_y_vel = P5.vel(N_in).dot(N_in.y)   #pos_from(P0).dot(N_in.y)
feet_y_vel = diff_to_symb_expr(feet_y_vel)
feet_y_vel
```

```python
cm_pos = m0*CM0.pos_from(P0)
cm_pos += m1*CM1.pos_from(P0)
cm_pos += m2*CM2.pos_from(P0)
cm_pos += m3*CM3.pos_from(P0)
cm_pos += m4*CM4.pos_from(P0)
cm_pos = cm_pos/(m0+m1+m2+m3+m4)

sys_CM = P0.locatenew('Sys_CM', cm_pos)
sys_CM_x = simplify(sys_CM.pos_from(P0).dot(N_in.x))
sys_CM_y = simplify(sys_CM.pos_from(P0).dot(N_in.y))
```

```python
sym_x = dynamicsymbols('q_0:5')
sym_x = sym_x + [ii.diff() for ii in sym_x]
sym_x = [diff_to_symb(ii) for ii in sym_x]
sym_params = list(symbols('I_0:5 d_0:5 g l_0:2 l_3 m_0:5'))
sym_add_params = [symbols('l_4'),]
sym_vars = sym_x + sym_params + sym_add_params

print(len(sym_vars), sym_vars)
```

```python
cas_x_args = cas.MX.sym("x", len(sym_x))
cas_params = cas.MX.sym("p", len(sym_params))
cas_add_params = cas.MX.sym("p_add", len(sym_add_params))
cas_all_vars = [cas_x_args[ii] for ii in range(len(sym_x))]
cas_all_vars += [cas_params[ii] for ii in range(len(sym_params))]
cas_all_vars += [cas_add_params[ii] for ii in range(len(sym_add_params))]
print(len(cas_all_vars), cas_all_vars)
```

```python
_cas_expr_temp_x = sympy2casadi(feet_x, sym_vars, cas_all_vars)
feet_x_cas = cas.Function(
        "Feet_x",
        [cas_x_args, cas_params, cas_add_params],
        [_cas_expr_temp_x,],
        ["x", "params", "additional_params"],
        ["feet_x_position"],
    )
```

```python
_cas_expr_temp_y = sympy2casadi(feet_y, sym_vars, cas_all_vars)
feet_y_cas = cas.Function(
        "Feet_y",
        [cas_x_args, cas_params, cas_add_params],
        [_cas_expr_temp_y,],
        ["x", "params", "additional_params"],
        ["feet_y_position"],
    )
```

```python
_cas_expr_temp_y_vel = sympy2casadi(feet_y_vel, sym_vars, cas_all_vars)
feet_y_vel_cas = cas.Function(
        "Feet_y_vel",
        [cas_x_args, cas_params, cas_add_params],
        [_cas_expr_temp_y_vel,],
        ["x", "params", "additional_params"],
        ["feet_y_speed"],
    )
```

```python
def simetric_cond_casadi(n = 5):
    x1 = cas.MX.sym('x_1', 2*n)
    x2 = cas.MX.sym('x_2', 2*n)
    cond = [x1[ii] - x2[n-1-ii] for ii in range(n)]
    cas_funcs = cas.horzcat(*cond)
    return cas.Function(
        "Sim_cond",
        [x1, x2],
        [cas_funcs,],
        ["x_1", "x2"],
        ["residue"],
    )
```

```python
simetric_5_links = simetric_cond_casadi(5)
```

Creating and simplifying symbolically the expressions of the heel impact may require some time, but it alows for a faster problem formulation later on.

```python
bodies = [body0, body1, body2, body3, body4]
points_right = [P0, P1, P2, P2, P4]
points_left = [P5, P4, P2, P2, P1]
subs_key = list(zip(dynamicsymbols('q_0:5'),dynamicsymbols('q_p_0:5')))

impact_eqs = []
for ii in range(5):
    print('calculating eq', ii)
    print('\tleft side')
    left_side = angular_momentum(points_left[ii], N_in, *bodies[:5-ii]).dot(N_in.z)
    left_side = simplify(left_side)
    print('\tright side')
    right_side = angular_momentum(points_right[ii], N_in, *bodies[ii:]).dot(N_in.z)
    right_side = simplify(right_side).subs(subs_key)
    impact_eqs.append(left_side-right_side)
#impact_eqs
```

```python
def impact_cond_casadi(eqs, x1_sym, x2_sym, sym_params, sym_add_params):
    x1_sym = [diff_to_symb(ii) for ii in x1_sym]
    x2_sym = [diff_to_symb(ii) for ii in x2_sym]
    eqs = [diff_to_symb_expr(ii) for ii in eqs]

    all_vars = x1_sym + x2_sym + sym_params + sym_add_params
    n = len(x1_sym)
    cas_x1 = cas.MX.sym('x_1', n)
    cas_x2 = cas.MX.sym('x_2', n)
    cas_params = cas.MX.sym("p", len(sym_params))
    cas_add_params = cas.MX.sym("p_add", len(sym_add_params))
    cas_all_vars = [cas_x1[ii] for ii in range(n)]
    cas_all_vars += [cas_x2[ii] for ii in range(n)]
    cas_all_vars += [cas_params[ii] for ii in range(len(sym_params))]
    cas_all_vars += [cas_add_params[ii] for ii in range(len(sym_add_params))]

    cas_funcs = []
    for function in eqs:
        cas_funcs.append(sympy2casadi(function, all_vars, cas_all_vars))
    cas_funcs = cas.horzcat(*cas_funcs)
    return cas.Function(
        "Sim_cond",
        [cas_x1, cas_x2, cas_params, cas_add_params],
        [cas_funcs,],
        ["x_1", "x2", 'params', 'additional_params'],
        ["residue"],
    )
```

```python
sym_x = dynamicsymbols('q_0:5')
sym_x = sym_x + [ii.diff() for ii in sym_x]
subs_key = list(zip(dynamicsymbols('q_0:5'),dynamicsymbols('q_p_0:5')))
sym_x_2 = [ii.subs(subs_key) for ii in sym_x]
impact_cond_cas_f = impact_cond_casadi(impact_eqs, sym_x,  sym_x_2, sym_params, sym_add_params)
```

```python
sys_cm_np = lambdify([sym_x, sym_params], [sys_CM_x, sys_CM_y],'numpy')
```

```python
ang_mom_p0 = angular_momentum(P0, N_in, *bodies).dot(N_in.z)
ang_mom_p0_np = lambdify([sym_x, sym_params], ang_mom_p0,'numpy')
```

```python
ang_mom_p5 = angular_momentum(P5, N_in, *bodies).dot(N_in.z)
ang_mom_p5_np = lambdify([sym_x, sym_params, sym_add_params], ang_mom_p5,'numpy')
```

```python
P5_static = P5.locatenew('P5_static', 0 * N_in.y)
P5_static.set_vel(N_in, 0 * N_in.y)
```

```python
ang_mom_p5_static = angular_momentum(P5_static, N_in, *bodies).dot(N_in.z)
ang_mom_p5_static_np = lambdify([sym_x, sym_params, sym_add_params], ang_mom_p5_static,'numpy')
```

```python
angular_momentum(P0, N_in, bodies[0]).dot(N_in.z)
```

```python
system_energy = potential_energy(*bodies) + kinetic_energy(N_in, *bodies)
```

```python
system_energy_np = lambdify([sym_x, sym_params], system_energy,'numpy')
```

```python
mass_matrix_np = lambdify([sym_x, sym_params], LM_small.mass_matrix,'numpy')
```

```python
sym_u = symbols('u_:5')
F_impl_np = lambdify([sym_x, sym_u, sym_params], LM_small.forcing,'numpy')
```

### Scheme definitions

Each scheme is defined here as a function that must be equal to zero at each interval.
Note that functions that contain "mod" in the name are those we define as "second order",
and use separate conditions for q and v.

Note that we will operate this problem wihout combining the scheme equations $F(q_k, q_{k+1}, q'_k, q'_{k+1} q''_k, q''_{k+1}, SchemeParams) = 0$ and the dynamics equations $H(q, q', q'', u, params) = 0$ imposed at the collocation points. This approach allows us to solve this problem without inverting the mass matrix.

If you wish to define your own schemes, do it here.

Be careful to respect the function structure:

    restriction(x, x_n, a, a_n, dt, scheme_params) = 0

```python
from optibot.schemes import index_div
from copy import copy

def euler_accel_restr(x, x_n, a, a_n, dt, scheme_params):
    first_ind, last_ind = index_div(x)
    x_d = copy(x)
    x_d[first_ind] = x[last_ind]
    x_d[last_ind] = a
    return x_n - (x + dt * x_d)

def trapz_accel_restr(x, x_n, a, a_n, dt, scheme_params):
    first_ind, last_ind = index_div(x)
    x_d = copy(x)
    x_d[first_ind] = x[last_ind]
    x_d[last_ind] = a
    x_d_n = copy(x)
    x_d_n[first_ind] = x_n[last_ind]
    x_d_n[last_ind] = a_n
    return x_n - (x + dt / 2 * (x_d + x_d_n))

def trapz_mod_accel_restr(x, x_n, a, a_n, dt, scheme_params):
    res = copy(x)
    first_ind, last_ind = index_div(x)
    res[last_ind] = x[last_ind] + dt / 2 * (a + a_n)
    res[first_ind] = x[first_ind] + dt * x[last_ind] + dt ** 2 / 6 * (a_n + 2 * a)
    return x_n - res

def hs_half_x(x, x_n, x_d, x_d_n, dt):
    x_c = (x + x_n) / 2 + dt / 8 * (x_d - x_d_n)
    return x_c

def hs_accel_restr(x, x_n, a, a_n, dt, scheme_params):
    a_c = scheme_params
    first_ind, last_ind = index_div(x)
    x_d = copy(x)
    x_d[first_ind] = x[last_ind]
    x_d[last_ind] = a

    x_d_n = copy(x)
    x_d_n[first_ind] = x_n[last_ind]
    x_d_n[last_ind] = a_n

    x_c = hs_half_x(x, x_n, x_d, x_d_n, dt)
    x_d_c = copy(x)
    x_d_c[first_ind] = x_c[last_ind]
    x_d_c[last_ind] = a_c
    return x + dt / 6 * (x_d + 4 * x_d_c + x_d_n) - x_n

def hs_mod_half_x(x, x_n, a, a_n, dt):
    x_c = copy(x)
    first_ind, last_ind = index_div(x)
    q = x[first_ind]
    v = x[last_ind]
    q_n = x_n[first_ind]
    v_n = x_n[last_ind]
    q_c = q + dt / 32 * (13 * v + 3 * v_n) + dt**2 / 192 * (11 * a - 5 * a_n)
    v_c = (v + v_n) / 2 + dt / 8 * (a - a_n)
    x_c[first_ind] = q_c
    x_c[last_ind] = v_c
    return x_c

def hs_mod_accel_restr(x, x_n, a, a_n, dt, scheme_params):
    a_c = scheme_params.T
    res = copy(x)
    first_ind, last_ind = index_div(x)
    q = x[first_ind]
    v = x[last_ind]
    res[last_ind] = v + dt / 6 * (a + 4 * a_c + a_n)
    res[first_ind] = q + dt * v + dt ** 2 / 6 * (a + 2 * a_c)
    return x_n - res
```

### Casadi optimization

We have generated the system equations symbolicaly. Now, we translate them to CasADi objects in order to perform the optimization.

```python
from optibot.casadi import accelrestriction2casadi
#from optibot.schemes import (euler_accel_restr, trapz_accel_restr, trapz_mod_accel_restr,
#                             hs_mod_accel_restr, hs_accel_restr, hs_half_x)
```

```python
#Numerical values of the paramenters

I_0_n, I_1_n, I_2_n, I_3_n, I_4_n = 0.93, 1.08, 2.22, 1.08, 0.93
d_0_n, d_1_n, d_2_n, d_3_n, d_4_n = 0.128, 0.163, 0.2, 0.163, 0.128
g_n = 9.81
l_0_n, l_1_n, l_2_n, l_3_n, l_4_n = 0.4, 0.4, 0.625, 0.4, 0.4
m_0_n, m_1_n, m_2_n, m_3_n, m_4_n = 3.2, 6.8, 20, 6.8, 3.2
params = [
    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,
    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,
    g_n,
    l_0_n, l_1_n, l_3_n,
    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n
]
additional_params = [l_4_n,]
```

```python
opti = cas.Opti()
p_opts = {}#{"expand":True,'ipopt.print_level':0, 'print_time':0}
s_opts = {}#{"max_iter": 10000, 'tol': 1e-26}#, 'linear_solver' : "MA27"}
opti.solver("ipopt",p_opts,
                    s_opts)
```

```python
N = 25
X = opti.variable(N+1,10)
X_dot = opti.variable(N+1,10)
U = opti.variable(N+1,5)
U_c = opti.variable(N,5)
X_c = opti.variable(N,10)
X_dot_c = opti.variable(N,10)
```

```python
T = opti.parameter()
u_m = opti.parameter()
Params_opti = opti.parameter(len(params))
Add_params_opti = opti.parameter(len(additional_params))
D = opti.parameter()
```

```python
# Definition of the cost function

#cost = cas.sum2((cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/N)
cost = cas.sum2((4*cas.sum1(U_c[:,:]**2) + cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/(3*N))
#cost = cas.sum2(cas.sum1(U**2))
opti.minimize(cost)
```

```python
#Periodic gait constraint:
opti.subject_to(simetric_5_links(X[0,:], X[-1,:]) == 0)
opti.subject_to(impact_cond_cas_f(X[-1,:], X[0,:], Params_opti, Add_params_opti) == 0)
```

```python
#Step size constraint:
opti.subject_to(feet_x_cas(X[-1,:], Params_opti, Add_params_opti) == D)
opti.subject_to(feet_y_cas(X[-1,:], Params_opti, Add_params_opti) == 0)
```

```python
#Small Feet Conditions:
opti.subject_to(U[:,0] == 0)
opti.subject_to(U_c[:,0] == 0)
opti.subject_to(feet_y_vel_cas(X[0,:], Params_opti, Add_params_opti)>0)
opti.subject_to(feet_y_vel_cas(X[-1,:], Params_opti, Add_params_opti)<0)
```

```python
#Feet over ground Restrictions:
for ii in range(1,N):
    opti.subject_to(feet_y_cas(X[ii,:], Params_opti, Add_params_opti) > 0)
```

```python
#Dynamics Constraints:
for ii in range(N+1):
    opti.subject_to(imp_dyn_x_f_cas(X[ii,:], X_dot[ii,:], U[ii,:], [], Params_opti) == 0)
for ii in range(N):
    opti.subject_to(X_c[ii,:] == hs_mod_half_x(X[ii,:], X[ii+1,:], X_dot[ii,5:], X_dot[ii+1,5:], T/N))
    opti.subject_to(imp_dyn_x_f_cas(X_c[ii,:], X_dot_c[ii,:], U_c[ii,:], [], Params_opti) == 0)
```

```python
#Scheme Constraints
#cas_accel_restr = accelrestriction2casadi(trapz_mod_accel_restr, 5)
cas_accel_restr = accelrestriction2casadi(hs_mod_accel_restr, 5, 5)
for ii in range(N):
    opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, X_dot_c[ii,5:]) == 0)
```

```python
opti.set_value(T, 0.7)#0.7
opti.set_value(D, 0.5)
```

```python
opti.set_value(Params_opti, params)
opti.set_value(Add_params_opti, additional_params)
```

```python
q_0_guess = np.array([-0.3, 0.7, 0, -0.5, -0.6])
q_1_guess = q_0_guess[::-1]
s_arr = np.linspace(0, 1, N+1)
q_guess = np.expand_dims(q_0_guess,0)+ np.expand_dims(s_arr,1)*np.expand_dims((q_1_guess - q_0_guess),0)

q_dot_guess = (q_1_guess - q_0_guess) * np.ones([N+1,1])/opti.value(T)
```

```python
opti.set_initial(X[:,:5], q_guess)
opti.set_initial(X[:,5:], q_dot_guess)

opti.set_initial(X_c[:,:5], (q_guess[:-1,:]+q_guess[1:,:])/2)
opti.set_initial(X_c[:,5:], q_dot_guess[:-1,:])

opti.set_initial(X_dot[:,:5], q_dot_guess)
opti.set_initial(X_dot[:,5:], 0)

opti.set_initial(X_dot_c[:,:5], q_dot_guess[:-1,:])
opti.set_initial(X_dot_c[:,5:], 0)

opti.set_initial(U, 0)
opti.set_initial(U_c, 0)
```

```python
sol = opti.solve()
```

```python
U_sol = sol.value(U)
U_c_sol = sol.value(U_c)
X_sol = sol.value(X)
X_c_sol = sol.value(X_c)
X_dot_sol = sol.value(X_dot)
X_dot_c_sol = sol.value(X_dot_c)
T_sol = sol.value(T)
T_sol_arr = np.linspace(0, T_sol, N+1)
T_c_arr = (T_sol_arr[:-1]+T_sol_arr[1:])/2
```

```python
plt.figure(figsize=[14,10])
labels= ['stance anckle', 'stance knee', 'stance hip', 'swing hip', 'swing knee']
for ii in range(5):
    plt.plot(T_sol_arr,U_sol[:,ii], marker = 'o', label = labels[ii] + ' u_k')
    plt.plot(T_c_arr,U_c_sol[:,ii], 'o', label = labels[ii] + ' u_c')
plt.grid()
plt.legend()
plt.title('u(t)')
```

```python
plt.figure(figsize=[14,10])
labels= ['stance tibia', 'stance femur', 'torso', 'swing femur', 'swing tibia']
for ii in range(5):
    plt.plot(T_sol_arr, X_sol[:,ii], marker = 'o', label = labels[ii] + ' q_k')
    plt.plot(T_c_arr,X_c_sol[:,ii], 'o', label = labels[ii] + ' q_c')
plt.grid()
plt.legend()
plt.title('q(t)')
```

```python
def chain_to_draw(x,params):
    [
    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,
    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,
    g_n,
    l_0_n, l_1_n, l_3_n,
    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n
    ] = params
    points_x = [0, ]
    points_y = [0, ]
    points_x.append(points_x[-1] - l_0_n*np.sin(x[0]))
    points_x.append(points_x[-1] - l_1_n*np.sin(x[1]))
    points_x.append(points_x[-1] - l_2_n*np.sin(x[2]))
    points_x.append(points_x[-2])
    points_x.append(points_x[-1] + l_3_n*np.sin(x[3]))
    points_x.append(points_x[-1] + l_4_n*np.sin(x[4]))

    points_y.append(points_y[-1] + l_0_n*np.cos(x[0]))
    points_y.append(points_y[-1] + l_1_n*np.cos(x[1]))
    points_y.append(points_y[-1] + l_2_n*np.cos(x[2]))
    points_y.append(points_y[-2])
    points_y.append(points_y[-1] - l_3_n*np.cos(x[3]))
    points_y.append(points_y[-1] - l_4_n*np.cos(x[4]))

    return points_x, points_y
```

```python
points_x, points_y = chain_to_draw(X_sol[0], params)
plt.figure(figsize=[15,15])
plt.grid()
for ii in range(0, N, 1):
    points_x, points_y = chain_to_draw(X_sol[ii], params)
    plt.plot(points_x, points_y, lw=1,  color = plt.cm.viridis(ii/N))

plt.gca().set_aspect('equal')
```

total_mass = m_0_n + m_1_n + m_2_n + m_3_n + m_4_n
ang_mom_arr = [ang_mom_p0_np(X_sol[ii,:],params) for ii in range(N+1)]
ang_mom_swing_foot_arr = [ang_mom_p5_np(X_sol[ii,:],params, additional_params) for ii in range(N+1)]
ang_mom_swing_foot_static_arr = [ang_mom_p5_static_np(X_sol[ii,:],params, additional_params) for ii in range(N+1)]
cm_torque_arr = [total_mass * -g_n * sys_cm_np(X_sol[ii,:], params)[0] for ii in range(N+1)]
ang_mom_arr_deriv = np.gradient(ang_mom_arr, T_sol_arr)

```python
from optibot.schemes import interpolated_array, interpolated_array_derivative
from optibot.analysis import dynamic_error_implicit
```

## Sistematic comparative of schemes for different values of N

Now let's solve the problem with different methods.

### Caution!

Executing the next cell may require some time!

```python
def q_init(N):
    q_0_guess = np.array([-0.3, 0.7, 0, -0.5, -0.6])
    q_1_guess = q_0_guess[::-1]
    s_arr = np.linspace(0, 1, N+1)
    q_guess = np.expand_dims(q_0_guess,0)+ np.expand_dims(s_arr,1)*np.expand_dims((q_1_guess - q_0_guess),0)
    q_dot_guess = (q_1_guess - q_0_guess) * np.ones([N+1,1])/opti.value(T)
    return q_guess, q_dot_guess
```

```python
import time
def chrono_solve(opti, solve_repetitions):
    cput0 = time.time()
    for ii in range(solve_repetitions):
        sol = opti.solve()
    cput1 = time.time()
    cpudt = (cput1-cput0)/solve_repetitions
    return sol, cpudt
```

```python
def casadi_biped(N = 25, scheme = "trapz", solve_repetitions = 1, t_end = 0.7, step_length = 0.5):
    opti = cas.Opti()
    p_opts = {"expand":True,'ipopt.print_level':0, 'print_time':0}
    s_opts = {"max_iter": 10000, 'tol': 1e-26}#, 'linear_solver' : "MA27"}
    opti.solver("ipopt",p_opts,
                        s_opts)

    restr_schemes = {
        'trapz': trapz_accel_restr,
        'trapz_mod' : trapz_mod_accel_restr,
        'hs': hs_accel_restr,
        'hs_mod': hs_mod_accel_restr,
        'hs_parab': hs_accel_restr,
        'hs_mod_parab': hs_mod_accel_restr
    }

    f_restr = restr_schemes[scheme]

    X = opti.variable(N+1,10)
    X_dot = opti.variable(N+1,10)
    U = opti.variable(N+1,5)
    if 'hs' in scheme:
        U_c = opti.variable(N,5)
        X_c = opti.variable(N,10)
        X_dot_c = opti.variable(N,10)

    T = opti.parameter()
    u_m = opti.parameter()
    Params_opti = opti.parameter(len(params))
    Add_params_opti = opti.parameter(len(additional_params))
    D = opti.parameter()

    # Cost
    if 'parab' in scheme:
        cost = cas.sum2((4*cas.sum1(U_c[:,:]**2) + cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/(3*N))
    else:
        cost = cas.sum2((cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/N)
    #cost = cas.sum2(cas.sum1(U**2))
    opti.minimize(cost)

    #Periodic gait constraint:
    opti.subject_to(simetric_5_links(X[0,:], X[-1,:]) == 0)
    opti.subject_to(impact_cond_cas_f(X[-1,:], X[0,:], Params_opti, Add_params_opti) == 0)

    #Step size constraint:
    opti.subject_to(feet_x_cas(X[-1,:], Params_opti, Add_params_opti) == D)
    opti.subject_to(feet_y_cas(X[-1,:], Params_opti, Add_params_opti) == 0)

    #Small Feet Conditions:
    opti.subject_to(U[:,0] == 0)
    opti.subject_to(feet_y_vel_cas(X[0,:], Params_opti, Add_params_opti)>0)
    opti.subject_to(feet_y_vel_cas(X[-1,:], Params_opti, Add_params_opti)<0)
    if 'hs' in scheme:
        opti.subject_to(U_c[:,0] == 0)

    #Feet over ground Restrictions:
    for ii in range(1,N):
        opti.subject_to(feet_y_cas(X[ii,:], Params_opti, Add_params_opti) > 0)

    #Dynamics Constraints:
    for ii in range(N+1):
        opti.subject_to(imp_dyn_x_f_cas(X[ii,:], X_dot[ii,:], U[ii,:], [], Params_opti) == 0)
    if 'hs' in scheme:
        for ii in range(N):
            opti.subject_to(X_c[ii,:] == hs_half_x(X[ii,:], X[ii+1,:], X_dot[ii,:], X_dot[ii+1,:], T/N))
            opti.subject_to(imp_dyn_x_f_cas(X_c[ii,:], X_dot_c[ii,:], U_c[ii,:], [], Params_opti) == 0)
        if 'parab' not in scheme:
            for ii in range(N):
                opti.subject_to(U_c[ii,:] == (U[ii,:]+U[ii+1,:])/2)

    #Scheme Constraints
    if 'hs' in scheme:
        cas_accel_restr = accelrestriction2casadi(f_restr, 5, 5)
        for ii in range(N):
            opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, X_dot_c[ii,5:]) == 0)
    else:
        cas_accel_restr = accelrestriction2casadi(f_restr, 5)
        for ii in range(N):
            opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, []) == 0)

    opti.set_value(T, t_end)#0.7
    opti.set_value(D, step_length)#0.5

    opti.set_value(Params_opti, params)
    opti.set_value(Add_params_opti, additional_params)

    q_guess, q_dot_guess = q_init(N)
    opti.set_initial(X[:,:5], q_guess)
    opti.set_initial(X[:,5:], q_dot_guess)

    opti.set_initial(X_dot[:,:5], q_dot_guess)
    opti.set_initial(X_dot[:,5:], 0)
    opti.set_initial(U, 0)

    if 'hs' in scheme:
        opti.set_initial(X_c[:,:5], (q_guess[:-1,:]+q_guess[1:,:])/2)
        opti.set_initial(X_c[:,5:], q_dot_guess[:-1,:])

        opti.set_initial(X_dot_c[:,:5], q_dot_guess[:-1,:])
        opti.set_initial(X_dot_c[:,5:], 0)
        opti.set_initial(U_c, 0)

    sol, cpudt = chrono_solve(opti, solve_repetitions)

    U_sol = sol.value(U)
    X_sol = sol.value(X)
    X_dot_sol = sol.value(X_dot)
    T_sol = sol.value(T)
    T_sol_arr = np.linspace(0, T_sol, N+1)
    T_c_arr = (T_sol_arr[:-1]+T_sol_arr[1:])/2
    cost_sol = sol.value(cost)
    if 'hs' in scheme:
        U_c_sol = sol.value(U_c)
        X_c_sol = sol.value(X_c)
        X_dot_c_sol = sol.value(X_dot_c)
    else:
        U_c_sol = None
        X_c_sol = None
        X_dot_c_sol = None

    return{
        'u':U_sol,
        'x':X_sol,
        'x_dot':X_dot_sol,
        't':T_sol,
        't_array':T_sol_arr,
        't_c_array': T_c_arr,
        'cpudt':cpudt,
        'u_c':U_c_sol,
        'x_c':X_c_sol,
        'x_dot_c':X_dot_c_sol,
        'cost':cost_sol
    }
```

```python
schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']
solve_repetitions = 3
N_arr = [20, 25, 30, 40, 50, 60]
results = {}

for scheme in schemes:
    key = scheme
    print('Problem:', key)
    results[key] = {'N_arr':N_arr}
    for N in N_arr:
        print(f'\tN = {N}')
        results[key][N] = casadi_biped(
            N = N,
            scheme = scheme,
            solve_repetitions = solve_repetitions,
            t_end = 0.7,
            step_length = 0.5)
```

### Calculating dynamic errors for each case

Caution! May take several seconds to run!

```python
schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']

n_graph = 2000 # A higher number here will provide more exact results but take longer to run
t_arr = np.linspace(0,0.7,n_graph)
for scheme in schemes:
    key = scheme
    if 'parab' in scheme:
        u_scheme = 'parab'
    else:
        u_scheme = 'lin'
    print('Problem:', key)
    N_arr = results[key]['N_arr']
    for N in N_arr:
        print(f'\tN = {N}')
        dyn_err_q, dyn_err_v, _, dyn_err_2 = dynamic_error_implicit(
            x_arr=results[key][N]['x'],
            u_arr=results[key][N]['u'],
            t_end=results[key][N]['t'],
            params = params,
            F_impl = F_impl_np,
            M = mass_matrix_np,
            scheme = scheme,
            u_scheme = u_scheme,
            scheme_params={'u_c':results[key][N]['u_c'],
                          'x_dot_c': results[key][N]['x_dot_c'],
                          'x_c': results[key][N]['x_c']},
            n_interp= n_graph)
        results[key][N]['dyn_err_q'] = dyn_err_q
        results[key][N]['dyn_err_v'] = dyn_err_v
        results[key][N]['dyn_err_2'] = dyn_err_2
```

```python
# Plot settings
plt.rcParams.update({'font.size': 15})
oct_fig_size = [10,6]
```

```python
schemes = ['hs_parab','hs_mod_parab', 'trapz', 'trapz_mod']
titles = ['Hermite Simpson','2nd order Hermite Simpson', 'Trapezoidal', '2nd order Trapezoidal']
colors = ['b', 'orange', 'g', 'r', 'purple']
n_int = len(t_arr)
N = 25
interv_n = (N * t_arr)/results[scheme][N]['t']
for kk in range(len(schemes)):
    scheme = schemes[kk]
    plt.figure(figsize=[14,8])
    for ii in range(5):
        cut_p = 0
        for ll in range(1,N+1):
            jj = np.searchsorted(interv_n, ll)
            plt.plot(t_arr[cut_p:jj],results[scheme][N]['dyn_err_q'][cut_p:jj,ii], '-', c = colors[ii], label = f'$q_{ii+1}$' if cut_p == 0 else None)
            cut_p = jj
    plt.plot(np.linspace(0,results[scheme][N]['t'],N+1), np.zeros(N+1), 'ok')
    plt.legend()
    plt.grid()
    if kk == 1:
        plt.ylim([-0.00001, 0.00001])
    elif kk == 3:
        plt.ylim([-0.001, 0.001])
    plt.title(r'First order dynamic error $\varepsilon^{[1]}_{q_i}$,'+f' {titles[kk]} scheme')
    plt.xlabel('Time(s)')
    plt.ylabel('Dynamic error $(rad/s)$')
    plt.tight_layout(pad = 0.0)
    sch_type = titles[kk].replace(' ','_')

    # If you are running the notebook locally and want to save the plots,
    # uncomment the next line
    #plt.savefig(f'5_link_First_Order_Dynamic_Error_{sch_type}_scheme.eps', format='eps')

```

```python
schemes = ['hs_parab','hs_mod_parab', 'trapz', 'trapz_mod']
titles = ['Hermite Simpson','2nd order Hermite Simpson', 'Trapezoidal', '2nd order Trapezoidal']
colors = ['b', 'orange', 'g', 'r', 'purple']
n_int = len(t_arr)
N = 25
interv_n = (N * t_arr)/results[scheme][N]['t']
for kk in range(len(schemes)):
    scheme = schemes[kk]
    plt.figure(figsize=[14,8])
    for ii in range(5):
        cut_p = 0
        for ll in range(1,N+1):
            jj = np.searchsorted(interv_n, ll)
            plt.plot(t_arr[cut_p:jj],results[scheme][N]['dyn_err_2'][cut_p:jj,ii], '-', c = colors[ii], label = f'$q_{ii+1}$' if cut_p == 0 else None)
            cut_p = jj
    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')
    if 'hs' in scheme:
        plt.plot(results[scheme][N]['t_c_array'], np.zeros(N), 'ow', markeredgecolor='b', label = 'collocation points')
        plt.ylim([-0.08, 0.08])
    else:
        plt.ylim([-1.75, 1.75])
    plt.legend()
    plt.grid()
    plt.title(r'Second order dynamic error $\varepsilon^{{[2]}}_{{q_i}}$,'+f' {titles[kk]} scheme')
    plt.xlabel('Time(s)')
    plt.ylabel('Dynamic error $(rad/s^2)$')
    plt.tight_layout(pad = 0.0)
    sch_type = titles[kk].replace(' ','_')

    # If you are running the notebook locally and want to save the plots,
    # uncomment the next line
    #plt.savefig(f'5_link_Second_Order_Dynamic_Error_{sch_type}_scheme.eps', format='eps')

```

```python
def arr_mod(x):
    x_1 = np.sum(x*x, axis=1)
    return np.sqrt(x_1)
def arr_sum(x):
    return np.sum(np.abs(x), axis = 1)
def arr_max(x):
    return np.max(np.abs(x), axis = 1)
```

```python
schemes = ['hs_mod_parab','hs_parab']#, 'trapz', 'trapz_mod']
titles = ['2nd order Hermite Simpson','Hermite Simpson']#, 'Trapezoidal', 'Modified Trapezoidal']
colors = ['b', 'orange', 'g', 'r', 'purple']
funcs = [arr_sum,]#arr_mod,  arr_max
#func_tittles = ['Module of', 'Sum of absolute', 'Maximum of absolute']
y_max_list = [0.12, 0.2, 0.09]
n_int = len(t_arr)
N = 25
interv_n = (N * t_arr)/results[scheme][N]['t']
for ii in range(1):
    plt.figure(figsize=oct_fig_size)
    for kk in [1,0]:
        scheme = schemes[kk]
        cut_p = 0
        for ll in range(1,N+1):
            jj = np.searchsorted(interv_n, ll)
            y_plot = funcs[ii](results[scheme][N]['dyn_err_2'])
            plt.plot(t_arr[cut_p:jj],y_plot[cut_p:jj], '-', c = f'C{kk}', label = titles[kk] if cut_p == 0 else None)
            cut_p = jj
    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')
    plt.plot(results[scheme][N]['t_c_array'], np.zeros(N), 'ow', markeredgecolor='k', label = 'collocation points')
    plt.legend()
    plt.grid()
    plt.ylim([-0.01,y_max_list[ii]])
    plt.title(r'Second order dynamic error $\varepsilon^{[2]}$,'+f' N = {N}')
    plt.xlabel('Time(s)')
    plt.ylabel('Dynamic error $(rad/s^2)$')
    plt.tight_layout(pad = 0.0)

    # If you are running the notebook locally and want to save the plots,
    # uncomment the next line
    #plt.savefig(f'5_link_HS_N{N}_second_order_dynamic_error.eps', format='eps')

```

```python
schemes = ['trapz', 'trapz_mod']
titles = ['Trapezoidal', '2nd order Trapezoidal']
funcs = [arr_sum,]#arr_mod,  arr_max
y_max_list = [0.12, 0.2, 0.09]
n_int = len(t_arr)
N = 50
interv_n = (N * t_arr)/results[scheme][N]['t']
for ii in range(1):
    plt.figure(figsize=oct_fig_size)
    for kk in range(2):
        scheme = schemes[kk]
        cut_p = 0
        for ll in range(1,N+1):
            jj = np.searchsorted(interv_n, ll)
            y_plot = funcs[ii](results[scheme][N]['dyn_err_2'])
            plt.plot(t_arr[cut_p:jj],y_plot[cut_p:jj], '-', c = f'C{kk+2}', label = titles[kk] if cut_p == 0 else None)
            cut_p = jj
    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')
    plt.legend()
    plt.grid()
    plt.title(r'Second order dynamic error $\varepsilon^{[2]}$,'+f' N = {N}')
    plt.xlabel('Time(s)')
    plt.ylabel('Dynamic error $(rad/s^2)$')
    plt.tight_layout(pad = 0.0)

    # If you are running the notebook locally and want to save the plots,
    # uncomment the next line
    #plt.savefig(f'5_link_Trapezoidal_N{N}_second_order_dynamic_error.eps', format='eps')

```

```python
def total_state_error(t_arr, dyn_err):
    errors = np.trapz(np.abs(dyn_err), t_arr, axis=0)
    return errors
```

```python
schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']
N_arr = [10,15,20,25,30,40,50,75,100,150]

t_arr = np.linspace(0,0.7,n_graph)
for scheme in schemes:
    key = scheme
    print('Problem:', key)
    N_arr = results[key]['N_arr']
    for N in N_arr:
        print(f'\tN = {N}')
        for letter in 'qv2':
            results[key][N][f'integ_dyn_err_{letter}']= total_state_error(t_arr, results[scheme][N][f'dyn_err_{letter}'])
            results[key][N][f'module_dyn_err_{letter}']= np.sqrt(np.dot(results[key][N][f'integ_dyn_err_{letter}'], results[key][N][f'integ_dyn_err_{letter}']))
            results[key][N][f'sum_dyn_err_{letter}'] = np.sum(results[key][N][f'integ_dyn_err_{letter}'])

```

```python
for scheme in schemes:
    key = scheme
    print('Problem:', key)
    N_arr = results[key]['N_arr']

    for letter in 'qv2':
        list_mod = []
        list_sum = []
        for N in N_arr:
            #print(f'\tN = {N}')
            list_mod.append(results[key][N][f'module_dyn_err_{letter}'])
            list_sum.append(results[key][N][f'sum_dyn_err_{letter}'])
        results[key][f'module_dyn_err_{letter}_array'] = np.array(list_mod)
        results[key][f'sum_dyn_err_{letter}_array'] = np.array(list_sum)
```

```python
# For each scheme, the number of collocation points can be obtained
for scheme in results.keys():
    if 'hs' in scheme:
        n_coll = np.array(results[scheme]['N_arr'])*2-1
        results[scheme]['N_coll_arr'] = n_coll
    else:
        results[scheme]['N_coll_arr'] = results[scheme]['N_arr']
```

```python
schemes = ['hs_mod_parab','hs_parab', 'trapz', 'trapz_mod']
titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']
plt.figure(figsize=oct_fig_size)
for ii in [2,3,1,0]:
    key = schemes[ii]
    plt.plot(results[key]['N_arr'], results[key][f'sum_dyn_err_2_array'], marker = 'o', c = f'C{ii}',label = titles[ii])
plt.grid()
plt.legend()
plt.yscale('log')
plt.title('Second order dynamic error $E^{[2]}$')
plt.xlabel('Number of intervals')
plt.ylabel('Dynamic error ($rad/s$)')
plt.tight_layout(pad = 0.0)

# If you are running the notebook locally and want to save the plots,
# uncomment the next line
#plt.savefig(f'5_link_Sum_second_order_dynamic_error_vs_interval_number.eps', format='eps')
```

```python
for scheme in schemes:
    key = scheme
    print('Problem:', key)
    N_arr = results[key]['N_arr']
    list_cpudt = []
    for N in N_arr:
        #print(f'\tN = {N}')
        list_cpudt.append(results[key][N]['cpudt'])
    results[key]['cpudt_array'] = np.array(list_cpudt)
```

```python
schemes = ['hs_mod_parab','hs_parab', 'trapz', 'trapz_mod']
titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']
plt.figure(figsize=oct_fig_size)
for ii in [2,3,1,0]:
    key = schemes[ii]
    plt.plot(results[key]['N_arr'], results[key][f'cpudt_array'], marker = 'o', c = f'C{ii}', label = titles[ii])
plt.grid()
plt.legend()
#plt.yscale('log')
plt.title('Optimization time')
plt.xlabel('Number of intervals')
plt.ylabel('Time (s)')
plt.tight_layout(pad = 0.0)

# If you are running the notebook locally and want to save the plots,
# uncomment the next line
#plt.savefig(f'5_link_optimization_vs_interval_number.eps', format='eps')
```

```python
# Here we print the data shown in Table III of the paper

for scheme in ['hs_mod_parab', 'hs_parab', 'trapz', 'trapz_mod']:
    key = scheme
    for N in [25,50]:#results[key]['N_arr']:
        print('scheme:', scheme, 'N:', N,'\n\ttime:', results[key][N][f'cpudt'],
              '\n\tErr 1:', results[key][N]['sum_dyn_err_q'], '\n\tErr 2:', results[key][N]['sum_dyn_err_2'])
```

## Animation

```python
from matplotlib import animation, rc
import matplotlib.patches as patches
from matplotlib.transforms import Affine2D
from IPython.display import HTML
import matplotlib
```

```python
matplotlib.rcParams['animation.embed_limit'] = 200
```

```python
def body_tray(X, params):
    res = []
    for ii in range(X.shape[0]):
        res.append(list(chain_to_draw(X[ii,:], params)))
    return np.array(res)
```

```python
def loop_body_tray(X, params):
    point_tray = body_tray(X, params)
    point_tray_loop = np.append(
        point_tray,
        np.expand_dims(
            np.array(list(chain_to_draw(X[0,[4,3,2,1,0,5,6,7,8,9]],params)))
            ,0),
        0)
    return point_tray_loop
```

```python
def mod_sum(iterable, start):
    for element in iterable:
        start += element
    return start
```

```python
def create_anim(X, U, params, n_loops = 1):
    [
    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,
    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,
    g_n,
    l_0_n, l_1_n, l_3_n,
    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n
    ] = params

    N = X.shape[0]
    fig, ax = plt.subplots()
    draw_width = 14
    draw_height = 14

    fig.set_dpi(72)
    fig.set_size_inches([draw_width,draw_height])
    ax.set_xlim(( -1, 1))
    ax.set_ylim(( -0.2, 1.8))

    body, = ax.plot([], [], lw=4, ms = 12, marker = 'o')
    trail, = ax.plot([], [], lw=1, color = 'k')
    old_trail, = ax.plot([], [], lw=1, color = 'k')
    next_trail, = ax.plot([], [], lw=1, color = 'k')

    point_tray = body_tray(X, params)
    point_tray_loop = loop_body_tray(X, params)
    #sys_cm_point, = ax.plot([], [], 'go', ms=12)
    #line_sys_cm, = ax.plot([], [], 'k:', lw=1)

    print_vars = [X[:,ii] for ii in range(5)]+[np.linspace(0, N-1, N, dtype=int)]
    print_var_names = [f'q_{ii}' for ii in range(5)]+['step']
    texts = []
    ii = 0.8
    for arr in print_vars:
        texts.append(ax.text(-0.8, ii, "", fontsize = 12))
        ii -= 0.2

    ax.grid()

    def init():
        body.set_data([], [])
        trail.set_data(point_tray_loop[0,0,-1], point_tray_loop[0,1,-1])
        old_trail.set_data(point_tray_loop[:,0,-1]-0.5, point_tray_loop[:,1,-1])
        #next_trail.set_data(point_tray_loop[:,0,-1]+0.5, point_tray_loop[:,1,-1])
        #sys_cm_point.set_data([], [])
        #line_sys_cm.set_data([], [])
        return (body,)

    def animate(i):
        margin_x = -0.25 + i * 0.5/N
        trail.set_data(point_tray_loop[0:i+1,0,-1], point_tray_loop[0:i+1,1,-1])
        #sys_cm_coords = sys_cm_np(X[i,:], params)
        #sys_cm_point.set_data(sys_cm_coords)
        #line_sys_cm.set_data([0, sys_cm_coords[0]], [0, sys_cm_coords[1]])

        ax.set_xlim(( -1+ margin_x, 1+ margin_x))
        points_x, points_y = point_tray[i,:,:]
        body.set_data(points_x, points_y)

        for ii in range(len(texts)):
            text = texts[ii]
            name = print_var_names[ii]
            arr = print_vars[ii]
            text.set_position((-0.9 + margin_x, 1.7 - 0.05*ii))
            if name == 'step':
                text.set_text("$step$ = " + str(arr[i]))
            else:
                text.set_text("$" + name + "$ = %.3f" % arr[i])
        return (body,)
    iterable_frames = mod_sum([[jj for jj in range(N)]for kk in range(n_loops)], start = [])
    anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=iterable_frames, interval=20,
                               blit=True)
    return anim

```

```python
anim = create_anim(results['hs_mod_parab'][25]['x'][:-1,:],results['hs_mod_parab'][25]['u'], params, 4)
```

```python
HTML(anim.to_jshtml())
```

```python
f = r"biped_animation.mp4"
writervideo = animation.FFMpegWriter(fps=25//0.7)
# If you are running the notebook locally and want to save the animation,
# uncomment the next line
#anim.save(f, writer=writervideo)
```

```python

```


# Bayes' Rule

## Motivating Example -- Part 1

Consider again a motivating example from the Introduction:

**A 40 year-old woman's mammogram comes back positive for a breast cancer tumor. What is the probability that the woman has breast cancer?**

This problem can be modeled as a *stochastic system*:
````{panels}
DEFINITION
^^^
stochastic system
: A stochastic system with one or more inputs and one or more outputs, for which the output(s) is(are) not a deterministic function of the output(s).
````

```{note}
One of the most challenging problems for students learning probability is taking a word problem, translating it into mathematics, identifying and finding the necessary information, and solving the problem.  There are several keys to being successful at this. The first two are:

**1.** Identify all of the phenomena (inputs, outputs, state) that may be random. Note that we should identify phenomenon that may be random **even if we are given a particular outcome, event, or realization occurred for that phenomenon**.
**2.** Introduce clearly-defined mathematical terminology for all of the random phenomena in the problem.
```

In this problem, we use randomness to model our lack of knowledge.  For example, for the woman under consideration, we do not know whether she has breast cancer or not. Thus, whether the woman actually has cancer is one random phenomenon. Although we are told that the mammogram came back positive, in general, it is a random phenomena that will depend on whether the woman being tested has breast cancer. These two are the only random phenomena in this problem, so we can not introduce notation.  Let
* $C$  is the event that a woman is diagnosed with having breast cancer in a given year
* $D$ is the event that a mammogram detects breast cancer in a woman in a given year

```{note}
The next steps in solving a probabiity word problem are:

**3.** Translate the probability that is being asked about from words to mathematics.
**4.** Use probability techniques to formulate a solution for the problem in terms of probabilities that are known or that can be found from reliable source.
```

Again, one of the hard things for people learning probability is to even write down a formula for the probability that the problem is asking about.  For instance, this problem states: "What is the probability that the woman has breast cancer?"  So, some readers may translate this probability to $$P(C)$$.

**That formula is not the probability that this problem is asking about!** The reason is that we need to interpret that sentence in terms of the other information given in the problem. In this case, the first sentence states "A 40 year-old woman's mammogram comes back positive for a breast cancer tumor". In other words, we are **given** that the mammogram came back positive for cancer, which in our mathematical notation is $D$.  Thus, the probability that this problem is seeking is $P(C|D)$, which is the *conditional* probability that the woman has cancer *given* that her mammogram detected cancer.

To begin to interpret this probability, the reader needs to understand medical tests are never perfect.
When a mammogram detects breast cancer, that does *not* necessarily mean that a woman has breast cancer. If a mammogram detects breast cancer and a woman does not actually have breast cancer, that is called a *false alarm*. The probability of a false alarm is $P(\overline{C}|D)$. You can find information on false alarm rates for mammograms on the Susan G. Komen Foundation's web pages at

https://www.komen.org/breast-cancer/screening/mammography/accuracy/

Based on the information on that page, let's let $P(D|\overline{C})=0.1$.

If a woman does have cancer, a mammogram should detect it most of the time. The probability of detecting an effect given it is actually present is called the *sensitivity* of the test. From the same Susan G. Komen Foundation web page, the sensitivity of a mammograms is $P(D|C)=0.87$.

```{note}
As a reminder, we should not expect $P(D| \overline{C})$ and $P(D|C)$ to add to 1 because they are computed using different conditional probability measures.
```

If a woman does have cancer, but the mammogram fails to detect it, it is called a *miss*. The probability of miss is $P(\overline{D}|C)=1- P(D|C)=0.13$

The probabilities $P(D|C)$, $P(\overline{D}|C)$, $P(D| \overline{C})$ and $P(\overline{D}|\overline{C})$ represent the probabilities of observing an effect given some true state (or input). These are called *likelihoods*:

````{panels}
DEFINITION
^^^
likelihood (discrete stochastic systems)
: Consider a stochastic system with a discrete set of possible input events $\{ A_0, A_i, \ldots \}$ and a discrete set of possible output events $\{B_0, B_1, \ldots\}$. Then the *likelihoods* are the conditional probabilities of the output events given the input events:
$$
P(B_j| A_i).
$$
````

Now we have a lot of different probabilities, but none of these are of the form $P(C|D)$. Instead, we only have probabilities like $P(D|C)$, $P(D|\overline{C})$, etc. (the likelihoods).

The probability $P(C|D)$ is the probability of the true state (or input) given the observed effect. It is called an *a posterior* or *posterior* probability, meaning that it is **after the observation or measurement**:

````{panels}
DEFINITION
^^^
*a posteriori* probability (discrete stochastic system)
: Consider a stochastic system with a discrete set of possible input events $\{ A_0, A_i, \ldots \}$ and a discrete set of possible output events $\{B_0, B_1, \ldots\}$. Then the *a posterior probabilities* are the conditional probabilities of the input events given the observed outputs:
$$
P(A_i | B_j).
$$

*A posteriori* probabilities are sometimes abbreviated APPs and sometimes called *posteriors*.
````

We need to develop a new formula to calculate an *a posteriori* probability from the likelihoods (and some additional information).

## Deriving Bayes' Rule

Bayes' Rule is a technique for expressing probabilities of the form $P(X|Y)$ in terms of probabilities of the form $P(Y|X)$ and some additional information. Thus, Bayes' rule can be used to find the *a posteriori probabilities* for our motivating problem.

We begin by expressing the desired probability using the definition of conditional probability:

$$
P(C|D) = \frac{P (C \cap D)}{P(D)}
$$

We can apply a chain rule to express $P(C \cap D)$ in terms of a likelihood:

$$
P(C|D) = \frac{P (D|C)P(C)}{P(D)}.
$$

Note however, that in addition to the likelihood, we also need $P(C)$. $P(C)$ is the probability of cancer without taking into account the mammogram, and is called an *a priori* probability:

````{panels}
DEFINITION
^^^
*a priori* probability (discrete stochastic system)
: Consider a stochastic system with a discrete set of possible input events $\{ A_0, A_i, \ldots \}$. Then the *a priori probabilities* are the probabilities of the input events before any output is observed or measured:

$$
P(A_i).
$$
*A priori* probabilities are sometimes called *priors* because they are the probabilities of the inputs *prior* to the output being observed or measured.
````

We are still left with $P(D)$, which we don't know. However, we can apply the Law of Total Probability to express $P(D)$ in terms of the likelihoods and *a priori* probabilities. Note that $C$ and $\overline{C}$ form a partition of $S$, so

$$
P(D) = P(D|C) P(C) + P(D| \overline{C}) P(\overline{C}).
$$

Finally, we have our formula for the desired *a posteriori* probability,

$$
P(C|D) = \frac{P (D|C)P(C)}{P(D|C) P(C) + P\left(D| \overline{C}\right) P\left(\overline{C}\right)}.
$$
This last formula is a form of **Bayes' Rule**:

````{panels}
DEFINITION
^^^
Bayes' Rule (discrete stochastic system)
: Consider a stochastic system with a discrete set of possible input events $\{ A_0, A_i, \ldots \}$ and a discrete set of possible output events $\{B_0, B_1, \ldots\}$.  Then the *a posteriori* probabilities $P(A_i|B_j)$ can be written in terms of the likelihoods ($P(B_j|A_i)$) and the *a priori* probabilities ($P(A_i)$) as

$$
P(A_i|B_j) = \frac{P\left(B_j \left \vert A_i \right. \right) P\left( A_i \right)}
{ \sum_i P\left(B_j \left \vert A_i \right. \right) P\left( A_i \right)}
$$

Note that the form of the numerator and the summands in the denominator are the same.
````

## Motivating Example -- Part 2

Now that we have applied Bayes' Rule to our problem, we have a formula for $P(C|D)$ in terms of the likelihoods (which we know) and the *a priori* probabilities (which we do not yet know). The *a priori* probabilities for this problem are $P(C)$ and $P\left( \overline{C} \right)$. Consider $P(C)$ -- this is the probability that our 40-year old woman has breast cancer without knowing the outcome of the mammogram.

As with the other probabilities, we can estimate $P(C)$ from data on the Susan G. Komen foundation website at https://www.komen.org/breast-cancer/risk-factor/age/#fig2-1

Shown below is a figure from that page that shows the number of breast cancer cases per 100,000 people in the United States.

Based on this figure, $P(C) \approx 140/100,000 = 1.4 \times10^{-3}$, and $P\left( \overline{C} \right) = 1- P\left( C \right)$.

Now we have enough data to calculate the conditional probability that a 40 year-old woman has cancer given that she has a mammogram that detects breast cancer:

\begin{align}
P(C|D) &= \frac{P (D|C)P(C)}{P(D|C) P(C) + P\left(D| \overline{C}\right) P\left(\overline{C}\right)} \\
&=\frac{(0.87)(1.4 \times 10^{-3})}{ (0.87)(1.4 \times 10^{-3}) +  \left[ \left( 0.1 \right) \left(1- 1.4 \times 10^{-3} \right)  \right]} \\
\end{align}

We will use Python to actually perform the calculation:

```python
PC = 1.4e-3
PNotC = 1-PC
PD_C = 0.87
PD_NotC = 0.1

PC_D= ( PD_C*PC )/ \
      (PD_C*PC + PD_NotC * PNotC)

print("The probability that a woman in her 40s actually has cancer when she has a positive mammogram test")
print(f"is approximately {PC_D:.3f}")
```

Is this result surprising to you? The test only has a false positive rate of 10%; i.e., it only indicates cancer when cancer is not present 10% of the time?  This is an example of the **base rate fallacy**:

````{panels}
DEFINITION
^^^
base rate fallacy
: Many people have incorrect intuition about dependent phenomena when given both likelihoods and *a priori* probabilities. The tendency is to focus on the likelihoods, which describe the relation among the phenomena and neglect the *a prioris*, which give some **base rate** at which one of the phenomena occurs. This is especially problematic when the base rate is particularly close to 0 or 1.  The base rate fallacy can be often be avoided by applying Bayes' rule to determine the *a posteriori* probabilities.
````

For more examples on the base rate fallacy, see the Wikipedia page:

[Base rate fallacy -- Wikpedia](https://en.wikipedia.org/wiki/Base_rate_fallacy).

In this case, the probability that a woman in her 40s has cancer is very low, so most women getting a mammogram will not have cancer. Thus, even a small false alarm rate results in a large number of women for which cancer is incorrectly detected.

It is easiest to interpret this result by treating the probabilities as rates and determining the number of women for which cancer is detected and the number of women for which cancer is falsely detected. We will focus our interpretation on the USA, so let's start by getting some data on how many 40 year-old women are in the US.  Type "US population by age and gender" into Google. You can then select a source for your data -- I would suggest the US Census Bureau, which took me to this page:

[Age & Sex Tables](https://www.census.gov/topics/population/age-and-sex/data/tables.html)

By clicking on "All", I got to a page that provides the information we want for 2019:

[Age and Sex Composition in the United States: 2019](https://www.census.gov/data/tables/2019/demo/age-and-sex/2019-age-sex-composition.html)

If you download the Excel spreadsheet linked to "Table 1. Population by Age and Sex: 2019", you should see something like the following:

According to the US Census Bureau, there were approximately 9,956,000 women from age 40 to 44 (5 years) in 2019. For convenience, let's estimate that there are now approximately 2 million women of age 40 -- the exact number is not important in this analysis.

Then the number of 40 year-old women with cancer is approximately

$$
2 \times 10^{6} \cdot \frac{140}{100,000} = 2800
$$

The number without cancer is $2 \times 10^6 - 2800 = 1,997,200$

According to the Kaiser Family Foundation, approximately 72% of women over age 40 have had a mammogram in the past two years:

https://www.kff.org/womens-health-policy/state-indicator/mammogram-rate-for-women-40-years/

The percentage is probably lower for 40 year-olds, so let's estimate the probability that a 40 year-old woman has a mammogram in a year as 1/4. If this rate applies equally to women who are ultimately determined to have cancer and those who do not, then the number of 40 year-old women that have a mammogram in a year is approximately 700 and 499,300 for those who ultimately have breast cancer and those who do not.

The number of women that have cancer and that will have cancer detected by a mammogram is approximately

$$
0.87(700) = 609
$$

The number of women that do not have cancer and that will have cancer detected by a mammogram is approximately

$$
0.1(499,300) = 49,930.
$$
Are you surprised? For every 40 year-old woman who has breast cancer correctly detected by a mammogram, there are almost 82 who have breast cancer **incorrectly detected** by a mammogram.

The total number of women for which breast cancer is detected is
$$
609+49,930= 50,539.
$$

Thus, the proportion of women for whom a mammogram detects breast cancer that actually have breast cancer is

$$
\frac{609}{50,539} \approx 0.012,
$$
which is the same result we got before. Using explicit numbers makes this easier to understand, but the important part of the math is exactly the same.

## Practice

Use the information from this section to answer the following question regarding mammograms and breast cancer:

```python
from jupyterquiz import display_quiz
git_path="https://raw.githubusercontent.com/jmshea/Foundations-of-Data-Science-with-Python/main/"

#display_quiz("quiz/bayes-rule.json")
display_quiz(git_path + "07-bayesian-methods/quiz/bayes-rule.json")
```

<div id="hnhBOlMKtyXm" data-shufflequestions="False"
               data-shuffleanswers="True"
               data-numquestions="1000000"> <style>:root {
    --medium-slate-blue: #6f78ffff;
    --orange-pantone: #f75c03ff;
    --russian-violet: #392061ff;
    --maximum-yellow-red: #ffc857ff;
    --viridian-green: #119da4ff;
    --incorrect-red: #c80202;
    --correct-green: #009113;
}

.Quiz {
    max-width: 600px;
    margin-top: 15px;
    margin-left: auto;
    margin-right: auto;
    margin-bottom: 15px;
    padding-bottom: 4px;
    padding-top: 4px;
    line-height: 1.1;
    font-size: 16pt;
}

.QuizCode {
    font-size: 14pt;
    margin-top: 10px;
    margin-left: 20px;
    margin-right: 20px;
}

.QuizCode>pre {
    padding: 4px;
}

.Answer {
    margin: 10px 0;
    display: grid;
    grid-template-columns: auto auto;
    grid-gap: 10px;
}

.Feedback {
    font-size: 16pt;
    text-align: center;
    min-height: 2em;
}

.Input {
    align: left;
    font-size: 20pt;
}

.Input-text {
    display: block;
    margin: 10px;
    color: inherit;
    width: 140px;
    background-color: #c0c0c0;
    color: #fff;
    padding: 5px;
    padding-left: 10px;
    font-family: inherit;
    font-size: 20px;
    font-weight: inherit;
    line-height: 20pt;
    border: none;
    border-radius: 0.2rem;
    transition: box-shadow 0.1s);
}

.Input-text:focus {
    outline: none;
    background-color: #c0c0c0;
    box-shadow: 0.6rem 0.8rem 1.4rem -0.5rem #999999;
}

.MCButton {
    background: #fafafa;
    border: 1px solid #eee;
    border-radius: 10px;
    padding: 10px;
    font-size: 16px;
    cursor: pointer;
    text-align: center;
}

.MCButton p {
    color: inherit;
}

.MultipleChoiceQn {
    padding: 10px;
    background: var(--medium-slate-blue);
    color: #fafafa;
    border-radius: 10px;
}

.ManyChoiceQn {
    padding: 10px;
    background: var(--orange-pantone);
    color: #fafafa;
    border-radius: 10px;
}

.NumericQn {
    padding: 10px;
    background: var(--russian-violet);
    color: #fafafa;
    border-radius: 10px;
}

.NumericQn p {
    color: inherit;
}

.InpLabel {
    line-height: 34px;
    float: left;
    margin-right: 10px;
    color: #101010;
    font-size: 15pt;
}

.incorrect {
    color: var(--incorrect-red);
}

.correct {
    color: var(--correct-green);
}

.correctButton {
    /*
    background: var(--correct-green);
   */
    animation: correct-anim 0.6s ease;
    animation-fill-mode: forwards;
    color: #fafafa;
    box-shadow: inset 0px 0px 5px #555555;
    outline: none;
}

.incorrectButton {
    animation: incorrect-anim 0.8s ease;
    animation-fill-mode: forwards;
    color: #fafafa;
    box-shadow: inset 0px 0px 5px #555555;
    outline: none;
}

@keyframes incorrect-anim {
    100% {
        background-color: var(--incorrect-red);
    }
}

@keyframes correct-anim {
    100% {
        background-color: var(--correct-green);
    }
}</style>

## Terminology Review

```python
from jupytercards import display_flashcards

#display_flashcards('flashcards/'+'bayes-rule.json')

github='https://raw.githubusercontent.com/jmshea/Foundations-of-Data-Science-with-Python/main/'
github+='07-bayesian-methods/flashcards/'
display_flashcards(github+'bayes-rule.json')

```

<style>
:root{
--asparagus: #87a878ff;
--terra-cotta: #e26d5aff;
--cyan-process: #5bc0ebff;
--dark-blue-gray: #666a86ff;
--snow: #fffafbff;
--rich-black-fogra-39: #090c08ff;
}

/* entire container, keeps perspective */
.flip-container {
	perspective: 1000px;
  touch-action: pinch-zoom pan-y;
}

.flip-container.flip .flipper {
	transform: rotateY(180deg)  translateX(-10%);
}

.flip-container.slide .flipper.frontcard {
    filter: brightness(90%);
    position:absolute;
    z-index: -20;
}

.flip-container.prepare .flipper.backcard {
    opacity:1;
    transform:  translateX(20px);
    transition: display 0s;
}

.flip-container.slide .flipper.backcard {
    filter: brightness(115%);
    height: 360px;
    opacity:1;
	  transform:  rotateY(3deg)  translateX(700px) translateZ(20px) scale(1, 1.05) rotate(3deg);
    z-index: 20;
}

.flip-container.slideback .flipper.backcard {
    filter: brightness(100%);
    opacity: 1;
	  transform: translateX(0px) translateZ(20px) rotateY(0deg);
    z-index: 10;
}

.flip-container.slideback .flipper.frontcard {
    filter: brightness(90%);
    transform: translateX(0px) translateZ(-20px);
    z-index: -10;
}

.flip-container, .front, .back {
	  height: 300px;
	  width: 640px;
}

.flipper.frontcard {
    position:absolute;
	  transform: translateX(0px) rotateY(0deg) translateZ(20px);
    z-index: 10;
}

.flipper.backcard {
    filter: brightness(90%);
    opacity:0;
    position:absolute;
    transform: translateX(0px) translateZ(-20px);
    z-index: -10;
}

/* flip speed goes here */
.flipper {
    cursor:pointer;
	  height: 200px;
    margin-left:20px;
	  position: absolute;
    top: 0;
	  transform-style: preserve-3d;
	  transition: 0.6s;
    width: 600px;
}

/* hide back of pane during swap */
.front, .back {
    -webkit-backface-visibility: hidden;
	  backface-visibility: hidden;
	  left: 0;
	  position: absolute;
	  top: 0;
}

/* front pane, placed above back */
.front {
	  backface-visibility: hidden;
    background: var(--asparagus);
    transform: rotateY(0deg);
	  z-index: 2;
}

/* back, initially hidden pane */
.back {
	  backface-visibility: hidden;
    background: var(--dark-blue-gray);
	  transform: rotateY(180deg);
}

.next {
    color: var(--rich-black-fogra-39);
    cursor:pointer;
    font-size: 16pt;
    left:600px;
    opacity:0.3;
    position:relative;
    top:20px;
}

.next.flipped {
    opacity:1;
}

.next.hide {
    opacity:0;
}

.flashcard{
    border-radius: 10px;
    color: #fafafa;
    display:block;
    padding:10px;

}
.flashcardtext{
    color: var(--snow);
    font-size: 20pt;
    margin-left:3%;
    position:absolute;
    text-align:center;
    top: 50%;
    transform: translate(0, -50%);
    width:90%;
  }

</style>
        <div style="height:40px"></div><div class="flip-container" id="TZOHHSKLoraF" onclick="flip(this)"></div><div style="height:40px"></div><div class="next" id="TZOHHSKLoraF-next" onclick="checkFlip('TZOHHSKLoraF')"> </div> <div style="height:40px"></div>

In the next section, we will explore how Bayes' Rule can be applied to a system with hidden state.

```python

```


# Realization of Non-Recursive Filters

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Fast Convolution

The straightforward convolution of two finite-length signals $x[k]$ and $h[k]$ is a numerically complex task. This has led to the development of various techniques with considerably lower complexity. The basic concept of the *fast convolution* is to exploit the correspondence between the convolution and the scalar multiplication in the frequency domain.

### Convolution of Finite-Length Signals

The convolution of a causal signal $x_L[k]$ of length $L$ with a causal impulse response $h_N[k]$ of length $N$ is given as

\begin{equation}
y[k] = x_L[k] * h_N[k] = \sum_{\kappa = 0}^{L-1} x_L[\kappa] \; h_N[k - \kappa] = \sum_{\kappa = 0}^{N-1} h_N[\kappa] \; x_L[k - \kappa]
\end{equation}

where $x_L[k] = 0$ for $k<0 \wedge k \geq L$ and $h_N[k] = 0$ for $k<0 \wedge k \geq N$. The resulting signal $y[k]$ is of finite length $M = N+L-1$. The computation of $y[k]$ for $k=0,1, \dots, M-1$ requires $M \cdot N$ multiplications and $M \cdot (N-1)$ additions. The computational complexity of the convolution is consequently [in the order of](https://en.wikipedia.org/wiki/Big_O_notation) $\mathcal{O}(M \cdot N)$. Discrete-time Fourier transformation (DTFT) of above relation yields

\begin{equation}
Y(e^{j \Omega}) = X_L(e^{j \Omega}) \cdot H_N(e^{j \Omega})
\end{equation}

Discarding the effort of transformation, the computationally complex convolution is replaced by a scalar multiplication with respect to the frequency $\Omega$. However, $\Omega$ is a continuous frequency variable which limits the numerical evaluation of this scalar multiplication. In practice, the DTFT is replaced by the discrete Fourier transformation (DFT). Two aspects have to be considered before a straightforward application of the DFT

1. The DFTs $X_L[\mu]$ and $H_N[\mu]$ are of length $L$ and $N$ respectively and cannot be multiplied straightforwardly

2. For $N = L$, the multiplication of the two spectra $X_L[\mu]$ and $H_L[\mu]$ would result in the [periodic/circular convolution](https://en.wikipedia.org/wiki/Circular_convolution) $x_L[k] \circledast h_L[k]$ due to the periodicity of the DFT. Since we aim at realizing the linear convolution $x_L[k] * h_N[k]$ with the DFT, special care has to be taken to avoid cyclic effects.

### Linear Convolution by Periodic Convolution

The periodic convolution of the two signals $x_L[k]$ and $h_N[k]$ is defined as

\begin{equation}
x_L[k] \circledast h_N[k] = \sum_{\kappa=0}^{M-1} \tilde{x}_M[k - \kappa] \; \tilde{h}_M[\kappa]
\end{equation}

where the periodic continuations $\tilde{x}_M[k]$ of $x_L[k]$ and $\tilde{h}_M[k]$ of $h_N[k]$ with period $M$ are given as

\begin{align}
\tilde{x}_M[k] &= \sum_{m = -\infty}^{\infty} x_L[m \cdot M + k] \\
\tilde{h}_M[k] &= \sum_{m = -\infty}^{\infty} h_N[m \cdot M + k]
\end{align}

The result of the circular convolution has a periodicity of $M$.

To compute the linear convolution by the periodic convolution one has to take care that the result of the linear convolution fits into one period of the periodic convolution. Hence, the periodicity has to be chosen as $M \geq N+L-1$. This can be achieved by zero-padding of $x_L[k]$ and $h_N[k]$ to a total length of $M$

\begin{align}
x_M[k] &= \begin{cases}
x_L[k] & \mathrm{for} \; k=0, 1, \dots, L-1 \\
0 & \mathrm{for} \; k=L, L+1, \dots, M-1
\end{cases}
\\
h_M[k] &= \begin{cases}
h_N[k] & \mathrm{for} \; k=0, 1, \dots, N-1 \\
0 & \mathrm{for} \; k=N, N+1, \dots, M-1
\end{cases}
\end{align}

This results in the desired equality of linear and periodic convolution

\begin{equation}
x_L[k] * h_N[k] = x_M[k] \circledast h_M[k]
\end{equation}

for $k = 0,1,\dots, M-1$ with $M = N+L-1$.

#### Example - Linear by periodic convolution

The following example computes the linear, periodic and linear by periodic convolution of a rectangular signal $x[k] = \text{rect}_L[k]$ of length $L$ with a triangular signal $h[k] = \Lambda_N[k]$ of length $N$.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

L = 32  # length of signal x[k]
N = 16  # length of signal h[k]
M = 16  # periodicity of periodic convolution

def periodic_summation(x, N):
    "Zero-padding to length N or periodic summation with period N."
    M = len(x)
    rows = int(np.ceil(M/N))

    if (M < int(N*rows)):
        x = np.pad(x, (0, int(N*rows-M)), 'constant')

    x = np.reshape(x, (rows, N))

    return np.sum(x, axis=0)

def periodic_convolve(x, y, P):
    "Periodic convolution of two signals x and y with period P."
    x = periodic_summation(x, P)
    h = periodic_summation(y, P)

    return np.array([np.dot(np.roll(x[::-1], k+1), h) for k in range(P)], float)

# generate signals
x = np.ones(L)
h = sig.triang(N)

# linear convolution
y1 = np.convolve(x, h, 'full')
# periodic convolution
y2 = periodic_convolve(x, h, M)
# linear convolution via periodic convolution
xp = np.append(x, np.zeros(N-1))
hp = np.append(h, np.zeros(L-1))
y3 = periodic_convolve(xp, hp, L+N-1)

# plot results
def plot_signal(x):
    plt.figure(figsize = (10, 3))
    plt.stem(x)
    plt.xlabel(r'$k$')
    plt.ylabel(r'$y[k]$')
    plt.axis([0, N+L, 0, 1.1*x.max()])

plot_signal(x)
plt.title('Signal $x[k]$')
plot_signal(y1)
plt.title('Linear convolution')
plot_signal(y2)
plt.title('Periodic convolution with period M = %d' %M)
plot_signal(y3)
plt.title('Linear convolution by periodic convolution');
```

**Exercise**

* Change the lengths `L`, `N` and `M` and check how the results for the different convolutions change

### The Fast Convolution

Using the above derived equality of the linear and periodic convolution one can express the linear convolution $y[k] = x_L[k] * h_N[k]$ by the DFT as

\begin{equation}
y[k] = \text{IDFT}_M \{ \; \text{DFT}_M\{ x_M[k] \} \cdot \text{DFT}_M\{ h_M[k] \} \; \}
\end{equation}

This operation requires three DFTs of length $M$ and $M$ complex multiplications. On first sight this does not seem to be an improvement, since one DFT/IDFT requires $M^2$ complex multiplications and $M \cdot (M-1)$ complex additions. The overall numerical complexity is hence in the order of $\mathcal{O}(M^2)$. The DFT can be realized efficiently by the [fast Fourier transformation](https://en.wikipedia.org/wiki/Fast_Fourier_transform) (FFT), which lowers the computational complexity to $\mathcal{O}(M \log_2 M)$. The resulting algorithm is known as *fast convolution* due to its computational efficiency.

The fast convolution algorithm is composed of the following steps

1. Zero-padding of the two input signals $x_L[k]$ and $h_N[k]$ to at least a total length of $M \geq N+L-1$

2. Computation of the DFTs $X[\mu]$ and $H[\mu]$ using a FFT of length $M$

3. Multiplication of the spectra $Y[\mu] = X[\mu] \cdot H[\mu]$

4. Inverse DFT of $Y[\mu]$ using an inverse FFT of length $M$

The overall complexity depends on the particular implementation of the FFT. Many FFTs are most efficient for lengths which are a power of two. It therefore can make sense, in terms of computational complexity, to choose $M$ as a power of two instead of the shortest possible length $N+L-1$. For real valued signals $x[k] \in \mathbb{R}$ and $h[k] \in \mathbb{R}$ the computational complexity can be reduced significantly by using a real valued FFT.

#### Example - Fast convolution

The implementation of the fast convolution algorithm is straightforward. Most implementations of the FFT include zero-padding to a given length $M$, e.g in `numpy` by `numpy.fft.fft(x, M)`. In the following example an implementation of the fast convolution is shown. For illustration the convolution of a rectangular signal $x[k] = \text{rect}_L[k]$ of length $L$ with a triangular signal $h[k] = \Lambda_N[k]$ of length $N$ is considered.

```python
L = 16  # length of signal x[k]
N = 16  # length of signal h[k]
M = N+L-1

# generate signals
x = np.ones(L)
h = sig.triang(N)

# linear convolution
y1 = np.convolve(x, h, 'full')
# fast convolution
y2 = np.fft.ifft(np.fft.fft(x, M) * np.fft.fft(h, M))

plt.figure(figsize=(10, 6))
plt.subplot(211)
plt.stem(y1)
plt.xlabel(r'$k$')
plt.ylabel(r'$y[k] = x_L[k] * h_N[k]$')
plt.title('Result of linear convolution')

plt.subplot(212)
plt.stem(y1)
plt.xlabel(r'$k$')
plt.ylabel(r'$y[k] = x_L[k] * h_N[k]$')
plt.title('Result of fast convolution')
plt.tight_layout()
```

#### Example - Numerical complexity

It was already argued that the numerical complexity of the fast convolution is considerably lower due to the usage of the FFT. The gain with respect to the convolution is evaluated in the following. In order to measure the execution times for both algorithms the `timeit` module is used. The algorithms are evaluated for the convolution of two random signals $x_L[k]$ and $h_N[k]$ of length $L=N=2^n$ for $n=0, 1, \dots, 16$.

```python
import timeit

n = np.arange(17)  # lengths = 2**n to evaluate
reps = 20  # number of repetitions for timeit

gain = np.zeros(len(n))
for N in n:
    length = 2**N
    # setup environment for timeit
    tsetup = 'import numpy as np; from numpy.fft import rfft, irfft; \
            x=np.random.randn(%d); h=np.random.randn(%d)' % (length, length)
    # direct convolution
    tc = timeit.timeit('np.convolve(x, x, mode="full")', setup=tsetup, number=reps)
    # fast convolution
    tf = timeit.timeit('irfft(rfft(x, %d) * rfft(h, %d))' % (2*length, 2*length), setup=tsetup, number=reps)
    # speedup by using the fast convolution
    gain[N] = tc/tf

# show the results
plt.figure(figsize = (15, 10))
plt.barh(n, gain, log=True)
plt.plot([1, 1], [-1, n[-1]+1], 'r-')
plt.yticks(n, 2**n)
plt.xlabel('Gain of fast convolution')
plt.ylabel('Length of signals')
plt.title('Comparison between direct/fast convolution')
plt.grid()
```

**Exercise**

* When is the fast convolution more efficient/faster than a direct convolution?
* Why is it slower below a given signal length?
* Is the trend of the gain as expected by the numerical complexity of the FFT?

Solution: The gain in execution time of a fast convolution over a direct implementation of the the convolution for different signal lengths depends heavily on the particular implementation and hardware used. The fast convolution in this example is faster for two signals having a length equal or larger than 1024 samples. Discarding the outliers and short lengths, the overall trend in the gain is approximately logarithmic as predicted above.

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Sascha Spors, Digital Signal Processing - Lecture notes featuring computational examples, 2016-2018*.


# Recursive Competitive Equilibria

Contents:

- [Recursive Competitive Equilibria](#Recursive-Competitive-Equilibria)
  - [Stochastic Growth Model](#Stochastic-Growth-Model)
  - [Transition Equations](#Transition-Equations)
  - [Steady State](#Steady-State)
  - [Policy Functions](#Policy-Functions)
  - [Simulations](#Simulations)

This lab includes:

(1) An overview of a basic stochastic growth model;

(2) Finding the transition equations of the model using dynamic programming;

(3) Solving for the policy functions of the model using value function iteration;

(4) Solving for the price system of the economy;

(5) Simulating the equilibrium paths of the economy under various initial conditions.

Load all necessary packages:

```julia
using Statistics, Plots, LinearAlgebra, Roots, QuantEcon, LaTeXStrings
```

---

## Stochastic Growth Model

### Refresher

Just to remind you of the nature of the model:

- The economy is populated by a **household** that consumes a good, sells labor, and trades claims to date $t+1$ consumption.

- On the production side, there is a firm that produces the consumption good. This is the **type I firm**.

- There is also a firm that transforms output into capital and earns revenues by renting out the accumulated capital to the type I firm. This is the **type II firm**.

### Additional Restrictions / Assumptions

Assume that our stochastic growth model has the exact same structure as the one presented in Chapter 12 of Ljungqvist and Sargent (covered in class yesterday).

We will apply some additional restrictions on the functional forms of the utility and production functions:

(1) The consumer's utility function is a function of consumption only: $u(C,1-N) = \theta \log(C) + (1-\theta) \log(1-N)$.

(2) The firm's production function is a function of capital only: $F(K,N) = F(K,N) = K^\alpha N^{1-\alpha}$.

### Equilibrium Objects

Recall that $X = \begin{bmatrix} K & A & s \end{bmatrix}$ is the vector of state variables, and $X'$ is equivalent to $X$ in the next period.

The following is a list of the objects that determine the **recursive competitive equilibrum**:

(1) The price system $\{ r(X), w(W), Q(X'\,|\,X) \}$;

(2) A perceived law of motion $K' = G(X)$ and an associated induced transition density $\hat{\Pi}(X'\,|\,X)$;

(3) A household value function $J(a,X)$ and decision rules $\sigma^c(a,X)$, $\sigma^n(a,X)$, and $\sigma^a(a,X;X')$.

### Obtaining Equilibrium Objects

The following is an outline of our approach to obtaining most of the equilibrium objects.

We largely defer to the results presented in section 12.9.2 of L&S.

**(1)** As our candidate for $G(\cdot)$ in $K' = G(X)$, we use the decision rule obtained through the planner's problem: $K' = \Omega^K(X)$.

**(2)** The price system is governed by the following equations:
\begin{align*}
r(X) &= A F_k(X) \, , \\
w(X) &= A F_n(X) \, , \\
Q(X'\,|\,X) &= \beta \Pi(X'\,|\,X) \frac{U_c(X')}{U_c(X)} \left[ A(s) F_K(X') + (1-\delta) \right] \, .
\end{align*}

**(3)** The policy functions $\sigma^c(a,X)$ and $\sigma^n(a,X)$ are equivalent to the policy functions $\Omega^C(X)$ and $\Omega^N(X)$ obtained through the planner's problem.

**(4)** The policy function for $a$, $\sigma^a(a,X;X')$, will also be pinned down by the equation $a = \left[ r(X) + (1-\delta) \right] K$ (section 12.9.1 of L&S).

### Planner's Problem

The Bellman equation for the planning problem is

\begin{align*}
    v(K,A,s) = \max_{C,N,K'} \left\{ u(C,1-N) + \beta \sum_{s'} \pi(s'\,|\,s) v(K',A',s') \right\} \, ,
\end{align*}

subject to

\begin{align*}
    K' + C &\leq A(s) F(K,N) + (1-\delta) K \, , \\
    A' &= A(s) \, .
\end{align*}

### Productivity Factor \& Events

Let $A_t = e^a_t$, where $a_t$ is the productivity shock and $A_t$ is the productivity factor.

We will assume that the productivity shock in our model is driven by the following autoregressive process:

\begin{align*}
a_t = \rho a_{t-1} + \gamma \varepsilon_t \, ,
\end{align*}

where $\varepsilon_t \sim N(0,1)$, and we have $\rho = 0.7$ and $\gamma = 1$.

Notice that $\mathbb{E}[a_t] = 0$, as well as $\mathbb{E}[A_t] = e^0 = 1$.

We will discretize this law of motion for $a$ and create an underlying Markov process with $s_t = \{1,2,\ldots,25\}$ and a stochastic matrix with entries corresponding to the nature of the given autoregressive process.

How we back out the Markov chain from the given process for the productivity factor is *not* important.

It *is* important to understand that the above process will give the necessary specification for the evolution of $A$ and $s_t$.

In the end we have $A' = A(s)$.

```julia
# Apply rouwenhorst to approximate productivity shock `productivity`
ρ = 0.7
σ = 0.01
productivity = rouwenhorst(25, ρ, σ, 0.0);
Π = (productivity.p) # save stochastic matrix
a = productivity.state_values # save productivity shock values
A = exp.(a); # save productivity factor shock values
```

### Final Additional Constraint

Just to make the analysis easier, assume that $\theta = 1$ and labor is supplies inelastically at $N=1$.

This effectively simplies our utility and production functions to:

\begin{align*}
u(C,1-N) &= u(C) = \log(C) \, , \text{ and } \\
F(K,N) &= F(K) = K^\alpha \, ,
\end{align*}

respectively.

---

## Transition Equations

To find the transition equations of the model, we start by solving the planner's problem.

### Simplifying the Planner's Bellman Equation

The Bellman equation:

\begin{align*}
    v(X) = \max_{C,K'} \left\{ \log(C) + \beta \mathbb{E} v(X') \right\} \, ,
\end{align*}

subject to

\begin{align*}
    K' + C &= A(s) K^\alpha + (1-\delta) K \, , \\
    A' &= A(s) \, .
\end{align*}

By setting the first flow constraint to equality, we are assuming an interior solution.

Use the first flow constraint to solve for $K'$ and substitute it into $X'$ in the Bellman equation:

\begin{align*}
    v(X) = \max_{C} \left\{ \log(C) + \beta \mathbb{E} v(A(s)K^\alpha + (1-\delta)K - C, A, s) \right\} \, .
\end{align*}

### FOC and Envelope Condition

The first-order condition (FOC) of the Bellman equation with respect to $C$ yields

\begin{align*}
    0 = \frac{1}{C} + \beta \mathbb{E} V'(X') (-1) \, ,
\end{align*}

which yields

\begin{align*}
    \frac{1}{C} = \beta \mathbb{E} V'(X') \, .
\end{align*}

The envelope condition with respect to $K$:

\begin{align*}
    V'(X) = \beta \mathbb{E}\left[ V'(X') \left( A(s) \alpha K^{\alpha-1} + 1 - \delta \right) \right] \, ,
\end{align*}

which is equivalent to

\begin{align*}
    V'(X) = \left( \alpha K^{\alpha-1} + 1 - \delta \right) \beta \mathbb{E}V'(X') \, .
\end{align*}

### Solving for the Euler Equation

Further substituting in the FOC into the above envelope condition yields

\begin{align*}
    V'(X) = \left( \alpha K^{\alpha-1} + 1 - \delta \right) \frac{1}{C} \, .
\end{align*}

Stepping the above equation for $V'(X)$ up by one period yields

\begin{align*}
    V'(X') = \left( \alpha K'^{\alpha-1} + 1 - \delta \right) \frac{1}{C'} \, .
\end{align*}

Substituting the above back into the FOC and solving for $C'$ yields the Euler equation:

\begin{align*}
    C' = C \beta \left( \alpha K'^{\alpha-1} + 1 - \delta \right) \, .
\end{align*}

With time indexes:

\begin{align*}
    C_{t+1} = C_t \beta \left( \alpha K_{t+1}^{\alpha-1} + 1 - \delta \right) \, .
\end{align*}

### All Transition Equations

We ended up with two transition equations.

The capital accumulation equation:

\begin{align}
\label{eq:capital_accumulation}
K_{t+1} = A(s_t) K_t^\alpha + (1-\delta)K_t - C_t \, .
\end{align}

The consumption Euler equation:

\begin{align}
\label{eq:euler}
C_{t+1} = C_t \beta \left( \alpha K_{t+1}^{\alpha-1} + 1 - \delta \right) \, .
\end{align}

These will come in handy when we try to simulate the model in its recursive competitive equilibrium.

---

## Steady State

We may use the given equations to solve for the following steady state values of the variables of interest:

\begin{align*}
\overline{K} &= \left( \frac{\alpha}{\beta^{-1} - 1 + \delta} \right)^{\frac{1}{1-\alpha}} \, , \\
\overline{Y} &= \overline{K}^\alpha \, , \\
\overline{C} &= \overline{Y} - \delta \overline{K} \, .
\end{align*}

Let $\alpha = 0.5$, $\delta = 0.1$, $\beta = 0.95$.

Show the above steady state values:

```julia
α = 0.5
δ = 0.1
β = 0.95
K̄ = ( α / ((1/β) - 1 + δ) )^(1/(1-α))
Ȳ = K̄^α
C̄ = Ȳ - δ*K̄
@show K̄, Ȳ, C̄;
```

    (K̄, Ȳ, C̄) = (10.731272294887045, 3.2758620689655182, 2.202734839476814)

---

## Policy Functions

### A Single Value Function Iteration

```julia
function bellmanMap(V, kgrid, A, Π, α, δ, β)

    N     = length(kgrid) #Number of gridpoints of capital
    S     = length(A) #Number of stochastic states

    V_new = zeros(S,N) #New Value function
    n_pol = zeros(Int,S,N) #New policy rule for grid points
    k_pol = zeros(S,N) #New policy rule for capital
    c_pol = zeros(S,N) #consumption rule

    obj   = zeros(N) #objective to be maximized

    EV    = Π*V #precompute expected value for speed

    for n in 1:N

        for s in 1:S

            for nprime in 1:N

                c = (1-δ)*kgrid[n] + A[s]*kgrid[n]^α - kgrid[nprime]

                if c <= 0
                    obj[nprime] = -Inf #penalty if consumption <0
                else
                    obj[nprime] = log(c) + β*EV[s,nprime] #otherwise compute objective
                end

            end

            #find optimal value and policy
            V_new[s,n], n_pol[s,n] = findmax(obj)
            k_pol[s,n]             = kgrid[n_pol[s,n]]
            c_pol[s,n]             = (1-δ)*kgrid[n] + A[s]*kgrid[n]^α - k_pol[s,n]
        end
    end

    return V_new, n_pol, k_pol, c_pol
end;
```

### Value Function Iteration (Until Convergence)

```julia
function solveBellman(V0, kgrid, A, Π, α, δ, β, ϵ=1e-6)
    diff = 1.
    V,n_pol,k_pol,c_pol = bellmanMap(V0,kgrid,A,Π,α,δ,β)
    while diff > ϵ
        V_new,n_pol,k_pol,c_pol = bellmanMap(V,kgrid,A,Π,α,δ,β)
        diff = norm(V_new-V,Inf)
        V = V_new
    end
    return V,n_pol,k_pol,c_pol
end;
```

### Solving for the Policy Functions

```julia
N = 99
kgrid = LinRange(K̄-10.0, K̄+10.0, N)
S = length(A)

# Solve the model
V,n_pol,k_pol,c_pol= solveBellman(zeros(S,N), kgrid, A, Π, α, δ, β);
```

```julia
plot(kgrid, k_pol[13,:], title = "Law of Motion for Capital", label = L"G(X)", lw=3)
xlabel!(L"k_t")
ylabel!(L"k_{t+1}")
```

```julia
plot(kgrid, c_pol[13,:], title = "Policy Function for Consumption", label = L"\sigma^c(X)", lw=3)
xlabel!(L"k_t")
ylabel!(L"c_t")
```

---

## Simulations

### RCE Simulation

```julia
function simulateRCE(n0, s0, T, n_pol, kgrid, Π, α, δ, β)
    k = zeros(T+1) # capital stock
    n = zeros(Int,T+1) #index of the capital stock
    c = zeros(T)
    y = zeros(T)
    r = zeros(T)
    Q = zeros(T)
    a = zeros(T)
    s = simulate_indices(MarkovChain(Π),T;init=s0)

    n[1] = n0
    k[1] = kgrid[n0]

    for t in 1:T
        c[t] = c_pol[s[t],n[t]]
        y[t] = A[s[t]]*k[t]^α
        n[t+1] = n_pol[s[t],n[t]] #get the policy rule for the index
        k[t+1] = kgrid[n[t+1]]
        r[t] = A[s[t]]*α*k[t]^(α-1)
        a[t] = (r[t] + 1 - δ)*k[t]

        if t < T
            Q[t] = β*Π[s[t],s[t+1]]*(c[t]/c_pol[s[t+1],n[t+1]])*(A[s[t+1]]*α*kgrid[n_pol[s[t+1],n[t+1]]]^(α-1)+1-δ)
        end
    end

    return k[1:T], c, y, r, Q, a
end
```

    simulateRCE (generic function with 1 method)

### Run the Simulation

```julia
T = 10000
n0 = 50
s0 = 13
k_sim, c_sim, y_sim, r_sim, Q_sim, a_sim = simulateRCE(n0, s0, T, n_pol, kgrid, Π, α, δ, β);
k = k_sim[3000:T]
c = c_sim[3000:T]
y = y_sim[3000:T]
r = r_sim[3000:T]
Q = Q_sim[3000:T]
a = a_sim[3000:T];
```

#### Output

```julia
plot(1:length(y), y, label = "Simulated Output", lw=0.3)
xlabel!("Time")
ylabel!("Output")
```

#### Capital

```julia
plot(1:length(k), k, label = "Simulated Capital", lw=1)
xlabel!("Time")
ylabel!("Capital")
```

#### Rental Rate of Capital

```julia
plot(1:length(r), r, label = "Simulated Rental Rate of Capital", lw=0.3)
xlabel!("Time")
ylabel!("Rental Rate of Capital")
```

#### Asset Price

```julia
plot(1:length(Q), Q, label = "Simulated Asset Price", lw=0.3)
xlabel!("Time")
ylabel!("Asset Price")
```

#### Asset Quantity

```julia
plot(1:length(a), a, label = "Simulated Asset Quantity", lw=1)
xlabel!("Time")
ylabel!("Asset Quantity")
```

---


# Two-level atomic system: 3 approaches #
In this notebook we will solve a two-level atomic system using three methods:
* the Liouville equations written by hand and solved analytically or using a solver from the library DifferentialEquations.jl ([see the nice tutorial by Simon Rochester](http://rochesterscientific.com/ADM/AtomicDensityMatrix/tutorial/TwoLevelSystem.html))
* the rate equations for Zeeman coherences (cf. [Blushs and Auzinsh, 2004](https://doi.org/10.1103/PhysRevA.69.063806))
* the Master Equation as implmented in QuantumOptics.jl (with some modifications needed, as we shall see).

```julia
using LinearAlgebra
using QuantumOptics
#import Pkg; Pkg.add("DifferentialEquations")
using DifferentialEquations
using Plots
plotly()
using Printf
using QuantumOptics.steadystate
```

```julia

```

## System description ##
As the name suggests, the two-level atom consists of a ground state $\left| 1 \right>$ and and excited state $\left| 2 \right>$. We will assume that the energy of the first level $\hbar \omega_1=0$ and the energy of the second level is $\hbar \omega=\Delta$. The two levels are coupled by a laser field with Rabi frequency $\Omega_R$. The decay from $\left| 1 \right>$ to $\left| 2 \right>$ is described by a decay rate $\Gamma_{21}$. Furthermore, atoms fly out of the beam at a rate $\gamma$ (fly-through relaxation), which affects both the ground and excited states. It is assumed that the ground state is repopulated at a rate $\lambda$ such that the overall population remains the same.

The system can be described by a density matrix $\rho$:

$$\rho = \begin{bmatrix}
\rho_{11} & \rho_{12}\\
\rho_{21} & \rho_{22}
\end{bmatrix}$$

The population of level $\left| 1 \right>$ is $\rho_{11}$, and the population of level $\left| 2 \right>$ is $\left| 1 \right>$ is $\rho_{11}$. The matrix elements $\rho_{12}$ and $\rho_{21}$ describe the coherences.

We will enter the parameter values once here: $\Omega_R$, $\Gamma$, $\gamma=\lambda$, and $\Delta$.

```julia
Ωᵣ=1.0+0.0im                # Rabi frequency
Γ=1.0+0.0im                 # decay rate from level 2 to 1.
γ=0.05+0.0im                # fly-through relaxation rate (assume same for all levels)
λ=γ                         # \lambda=\gamma Tr(\rho), and Tr(\rho)=1 for constant population.
Δ=0.1+0.0im                 # laser detuning.
p=[Ωᵣ,Γ,γ,Δ]                # paramter vector
ρ0 = [1.0+0.0im, 0.0+0.0im] # initial state in ground state (when needed)
tspan = (0.0,10.0)          # time span (for time-dependent problems)
```

    (0.0, 10.0)

## Liouville equations ##
The problem described above can be reprented by the following system of differential equations:
$$
\begin{array}{lcl}
\frac{\partial \rho_{11}}{\partial t} &=& \lambda - \gamma \rho_{11} -i\frac{\Omega_R}{2}(\rho_{12}-\rho_{21})+\Gamma \rho_{22} \\
\frac{\partial \rho_{12}}{\partial t} &=& \frac{1}{2}\left[ -\gamma \rho_{12} -(\gamma+\Gamma)\rho_{12}\right] -i\frac{\Omega_R}{2}(\rho_{11}-\rho_{22}) -i\Delta \rho_{12}\\
\frac{\partial \rho_{21}}{\partial t} &=& \frac{1}{2}\left[ -\gamma \rho_{21} -(\gamma+\Gamma)\rho_{21}\right] -i\frac{\Omega_R}{2}(-\rho_{11}+\rho_{22}) +i\Delta \rho_{21}\\
\frac{\partial \rho_{22}}{\partial t} &=& -i\frac{\Omega_R}{2}(-\rho_{12}+\rho{21})-(\gamma + \Gamma)\rho_{22}
\end{array}
$$

```julia
# We define the Liouville system
function Liouville!(dρ₂,ρ₂,p,t)
    Ωᵣ,Γ,γ,Δ = p
    dρ₂[1]=γ-γ*ρ₂[1]-1.0im*0.5*Ωᵣ*(ρ₂[2]-ρ₂[3])+Γ*ρ₂[4]
    dρ₂[2]=0.5*( -γ*ρ₂[2]-(γ+Γ)*ρ₂[2] ) - 1.0im*0.5*Ωᵣ*( ρ₂[1] - ρ₂[4] ) - 1.0im*Δ*ρ₂[2]
    dρ₂[3]=0.5*( -γ*ρ₂[3]-(γ+Γ)*ρ₂[3] ) - 1.0im*0.5*Ωᵣ*( -ρ₂[1] + ρ₂[4] ) + 1.0im*Δ*ρ₂[3]
    dρ₂[4]=-(γ+Γ)*ρ₂[4]-1.0im*0.5Ωᵣ*( -ρ₂[2] + ρ₂[3] )
end
```

    Liouville! (generic function with 1 method)

### Analytical solution of the Liouville equations ###
The Liouville equations for this system have an exact solution for the steady-state value of the density matrix.

```julia
# exact solution of Liouville system (from Rochester)
ΩR=Ωᵣ
ρ₁₁=((γ+Γ)*((2*γ+Γ)^2+4*Δ^2)+(2*γ+Γ)*ΩR^2)/(((γ+Γ)*((2*γ+Γ)^2+4*Δ^2))+2*(2*γ+Γ)*ΩR^2)
ρ₁₂=(0.0-1.0*im)*(γ+Γ)*(2*γ+Γ-2*(1.0+1.0im)*Δ)*ΩR/(((γ+Γ)*((2*γ+Γ)^2+4*Δ^2))+2*(2*γ+Γ)*ΩR^2)
ρ₂₁=(0.0+1.0*im)*(γ+Γ)*(2*γ+Γ+2*(1.0+1.0im)*Δ)*ΩR/(((γ+Γ)*((2*γ+Γ)^2+4*Δ^2))+2*(2*γ+Γ)*ΩR^2)
ρ₂₂=(2*γ+Γ)*ΩR^2/(((γ+Γ)*((2*γ+Γ)^2+4*Δ^2))+2*(2*γ+Γ)*ΩR^2)
println("ρ₁₁ = ", ρ₁₁)
println("ρ₁₂ = ", ρ₁₂)
println("ρ₂₁ = ", ρ₂₁)
println("ρ₂₂ = ", ρ₂₂)
println("Tr(ρ) = ",ρ₁₁+ρ₂₂)
```

    ρ₁₁ = 0.6868327402135233 + 0.0im
    ρ₁₂ = -0.0597864768683274 - 0.26903914590747335im
    ρ₂₁ = -0.0597864768683274 + 0.38861209964412813im
    ρ₂₂ = 0.31316725978647686 + 0.0im
    Tr(ρ) = 1.0 + 0.0im

### Solving the Liouville equations with DifferentialEquations.jl ###
We can also solve the Liouville equations using solvers from Differential Equations.

#### Steady-state problem ####
First, we will solve the steady-state problem using the solver SSRootfind().

```julia
ρ2₀ = [1.0+0.0im, 0.0+0.0im, 0.0+0.0im, 0.0+0.0im]     # initial conditions (not really needed for steady state)
probLiouville = SteadyStateProblem(Liouville!,ρ2₀,p)
sol = solve(probLiouville,SSRootfind())
```

    u: 4-element Vector{ComplexF64}:
        0.6868327402122938 - 1.0547776057540605e-16im
      -0.05978647686767886 - 0.328825622776546im
     -0.059786476867678774 + 0.328825622776546im
        0.3131672597871866 - 6.384439715179834e-17im

#### Time-dependent problem ####
However, we can also solve it as a time-dependent problem.

```julia
ρ2₀=[1.0+0.0im,0.0+0.0im,0.0+0.0im,0.0+0.0im]
tspan = (0.0,10.0)
probLiouville = ODEProblem(Liouville!,ρ2₀,tspan,p)
sol = solve(probLiouville)
plot(sol)
```

## Rate equations for Zeeman coherences ##
One can average over the coherence terms in the Liouville equations (cf. [Blushs and Auzinsh, 2004](https://doi.org/10.1103/PhysRevA.69.063806)). This procedure has the advantage or reducing the number of equations, but it is also useful as a way to account for phase fluctuations in the laser beam, which are reduced to a single parameter, the laser linewidth. Details can be found in Blushs and Auzinsh (2004). The result is:
$$\begin{array}{lcl}
    \frac{\partial \rho_{11}}{\partial t} &=& \frac{1}{4}(\Xi_{12}+\Xi_{21})\rho_{22} - \frac{1}{4}(\Xi_{12}+\Xi_{21})\rho_{11}  - \gamma \rho_{11} +\Gamma \rho_{22} + \lambda \\
\frac{\partial \rho_{22}}{\partial t} &=& \frac{1}{4}(\Xi_{21}+\Xi_{12})\rho_{11} - \frac{1}{4}(\Xi_{12}+\Xi_{21})\rho_{22} -(\gamma + \Gamma)\rho_{22}
\end{array}
$$
where,
$$\begin{array}{lcl}
\Xi_{12}=\frac{\Omega_R^2}{\left(\frac{\Gamma}{2}+\gamma + i\Delta \right)} \\
\Xi_{21}=\frac{\Omega_R^2}{\left(\frac{\Gamma}{2}+\gamma - i\Delta \right)}
\end{array}$$
and the other constants are as defined earlier.

```julia
#Divu limenu sistema. Linarda pieeja
function bliv!(dρ,ρ,p,t)
    Ωᵣ,Γ,γ,Δ = p
    Ξ₁₂ = Ωᵣ^2 /((Γ/2)+γ +1.0im*(Δ))
    Ξ₂₁ = Ωᵣ^2 /((Γ/2)+γ -1.0im*(Δ))
    dρ[1] = 0.25*( Ξ₁₂ + Ξ₂₁ )*ρ[2] - 0.25*( Ξ₂₁ + Ξ₁₂ )*ρ[1] + Γ*ρ[2] - γ*ρ[1] + γ
    dρ[2] = 0.25*( Ξ₂₁ + Ξ₁₂ )*ρ[1] - 0.25*( Ξ₁₂ + Ξ₂₁ )*ρ[2] - (Γ + γ)*ρ[2]
end
```

    bliv! (generic function with 1 method)

```julia
prob = SteadyStateProblem(bliv!,ρ0,p)
sol = solve(prob)
```

    u: 2-element Vector{ComplexF64}:
       0.686832740251305 + 0.0im
     0.31316725980611354 + 0.0im

## Formulating the problem using the Master Equation ##
The two-level problem can be described in a different way that some people may find more intuitive by writing the Hamiltonian as follows:
$$\begin{equation}
\hat{H}=\frac{\Omega_R}{2} \left( \left| 1 \right> \left< 2 \right|+ \left| 2 \right> \left< 1 \right| \right) - \Delta \left| 2 \right> \left< 2 \right|
\end{equation}
$$
Relaxation terms can be included using the Lindblad superoperator, for which we define Jump operators. The Jump operator for decay from level 1 to level 2 is (somewhat counterintuitively):
$$
\hat{J}=\sqrt(\Gamma_{21})\left| 1 \right> \left< 2 \right|
$$
Then the Lindblad superoperator is
$$
\hat{\mathcal{L}}[\rho]=J \rho J^{\dagger} - \frac{1}{2} \left( J^{\dagger}J \rho + \rho J^{\dagger}J \right)
$$

For our two-level system, fly-through relaxation introduces decoherence, which we can describe for each level as
$$
\hat{J}^{fly-through}_k=\sqrt(\gamma)\left| k \right> \left< k \right|.
$$

Then, the system dynamics are described by the equation
$$
\frac{\partial \rho}{\partial t}=-\frac{i}{\hbar}[H,\rho] + \sum_{i}\hat{\mathcal{L}}_i[\rho]
$$

The Library QuantumOptics.jl provides a function called liouvillian that transforms a system defined as above into a matrix *L* such that
$$
\frac{\partial \mathbf{\rho}}{\partial t} = \mathbf{L} \mathbf{\rho},
$$
where $\mathbf{\rho}$ is determined as a vector
$$
\rho = \begin{bmatrix}
\rho_{11} \\
\rho_{12} \\
\rho_{21} \\
\rho_{22}
\end{bmatrix}
$$
However, the transfer of atoms into and outof the laser beam by fly-through relaxation is not included in this way (only the decoherence introduced in this way), so we have to include it by hand by adding the appropriate terms to the Liouville matrix and by making the problem inhomogeneous, i.e., $\frac{\partial \rho}{\partial t}=-\lambda$, where $\lambda = \gamma$.

```julia
my2levelSystem=NLevelBasis(2)
σ₁₂=transition(my2levelSystem,1,2)   # Laser excitation |1><2|
σ₂₁=transition(my2levelSystem,2,1)   # |2><1|
Γ₂₁=transition(my2levelSystem,1,2)   # excited state relaxation |2><1|

proj2=transition(my2levelSystem,2,2) # |2><2|
proj1=transition(my2levelSystem,1,1)

J2=[sqrt(Γ)*Γ₂₁,
   (sqrt(γ))*proj1,
   (sqrt(γ))*proj2
  ]
H2=(ΩR/2)*(σ₁₂+σ₂₁) + Δ*proj2
L=liouvillian(H2,J2)
DenseL=Array(L.data)
A=DenseL

A[1,1]=A[1,1]-γ            # add terms for fly-through population loss by hand
A[4,4]=A[4,4]-γ            # add terms for fly-through population loss by hand
y=[-λ,0,0,0]               # inhomogeneous term (to account for repopulation of fly-through)
Ainv=inv(A)                # invert matrix
eigen_vectors=Ainv*y
```

    4-element Vector{ComplexF64}:
        0.6868327402135219 + 1.92906189243143e-17im
     -0.059786476868327304 - 0.32882562277580013im
     -0.059786476868327304 + 0.3288256227758001im
        0.3131672597864763 + 0.0im


# SymPy QUIZ

```python
import sympy as sp
```

```python
sp.init_printing()
x = sp.symbols('x')
```

## What is the denominator of the third term of the Taylor expansion for the equation

$$ \Large \frac{\sin{\left (x \right )}}{x^{4} + 27}$$

```python

```

```python

```

```python

```

## Answer the Canvas quiz `SymPyQuiz`