# Cart-pole swing-up 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 Cart-pole swing-up 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
```

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

```python
from functools import lru_cache
```

```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, l, t, g = symbols('m_0 m_1 l t g')
q0, q1 = dynamicsymbols('q_0 q_1')
```

```python
# Definition of the physics system

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

P0 = pN.locatenew('P0', q0 * N_in.x)
P0.set_vel(N_in, q0.diff(t) * N_in.x)
cart_part = Particle('CartPart', P0, m0)
cart_part.potential_energy = m0 * g * P0.pos_from(pN).dot(N_in.y)

N1 = N_in.orientnew('N1', 'Axis', [q1, N_in.z])
P1 = P0.locatenew('P1', -l*N1.y)
P1.set_vel(N_in, P1.pos_from(pN).dt(N_in))

pend_part = Particle('PendPart', P1, m1)
pend_part.potential_energy = m1 * g * P1.pos_from(pN).dot(N_in.y)
```

```python
#Computing the Lagrangian

Lag_simp = Lagrangian(N_in, cart_part, pend_part)
Lag_simp
```

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

u0, u1 = symbols('u_0, u_1')
FL = [(P0, u0 * N_in.x)]#, (N1, u1 * N_in.z)]
LM_small = SimpLagrangesMethod(Lag_simp, [q0, q1], forcelist=FL, frame=N_in)
```

```python
# Generating the dynamic equations

LM_small.form_lagranges_equations()
RHS_small = LM_small.rhs
RHS_small
```

### 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.

Schemes that contain "parab" in the name are versions of Hermite Simpson that allow
or $U_c$ to be a free parameter. It is passed to the function through the
"scheme_params" argument.

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

Be careful to respect the function structure: either

    restriction(x, x_n, u, u_n, F, dt, params) = 0
or

    restriction(x, x_n, u, u_n, F, dt, params, scheme_params) = 0

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

def euler_restr(x, x_n, u, u_n, F, dt, params):
    return x_n - (x + dt * F(x, u, params))

def trapz_restr(x, x_n, u, u_n, F, dt, params):
    f = F(x, u, params)
    f_n = F(x_n, u_n, params)
    return x_n - (x + dt / 2 * (f + f_n))

def trapz_mod_restr(x, x_n, u, u_n, F, dt, params):
    res = copy(x)
    first_ind, last_ind = index_div(x)
    q = x[first_ind]
    v = x[last_ind]
    f = F(x, u, params)[last_ind]
    f_n = F(x_n, u_n, params)[last_ind]
    res[last_ind] = v + dt / 2 * (f + f_n)
    res[first_ind] = q + dt * v + dt ** 2 / 6 * (f_n + 2 * f)
    return x_n - res

def hs_restr(x, x_n, u, u_n, F, dt, params):
    f = F(x, u, params)
    f_n = F(x_n, u_n, params)
    x_c = (x + x_n) / 2 + dt / 8 * (f - f_n)
    u_c = (u + u_n) / 2
    f_c = F(x_c, u_c, params)
    return x + dt / 6 * (f + 4 * f_c + f_n) - x_n

def hs_mod_restr(x, x_n, u, u_n, F, dt, params):
    x_c = copy(x)
    res = copy(x)
    first_ind, last_ind = index_div(x)
    f = F(x, u, params)[last_ind]
    f_n = F(x_n, u_n, params)[last_ind]
    q = x[first_ind]
    v = x[last_ind]
    q_n = x_n[first_ind]
    v_n = x_n[last_ind]
    u_c = (u + u_n) / 2
    q_c = q + dt / 32 * (13 * v + 3 * v_n) + dt**2 / 192 * (11 * f - 5 * f_n)
    v_c = (v + v_n) / 2 + dt / 8 * (f - f_n)
    x_c[first_ind] = q_c
    x_c[last_ind] = v_c
    f_c = F(x_c, u_c, params)[last_ind]
    res[last_ind] = v + dt / 6 * (f + 4 * f_c + f_n)
    res[first_ind] = q + dt * v + dt ** 2 / 6 * (f + 2 * f_c)
    return x_n - res

def hs_parab_restr(x, x_n, u, u_n, F, dt, params, scheme_params):
    f = F(x, u, params)
    f_n = F(x_n, u_n, params)
    x_c = (x + x_n) / 2 + dt / 8 * (f - f_n)
    u_c = scheme_params
    f_c = F(x_c, u_c, params)
    return x + dt / 6 * (f + 4 * f_c + f_n) - x_n

def hs_mod_parab_restr(x, x_n, u, u_n, F, dt, params, scheme_params):
    x_c = copy(x)
    res = copy(x)
    first_ind, last_ind = index_div(x)
    f = F(x, u, params)[last_ind]
    f_n = F(x_n, u_n, params)[last_ind]
    q = x[first_ind]
    v = x[last_ind]
    q_n = x_n[first_ind]
    v_n = x_n[last_ind]
    u_c = scheme_params
    q_c = q + dt / 32 * (13 * v + 3 * v_n) + dt**2 / 192 * (11 * f - 5 * f_n)
    v_c = (v + v_n) / 2 + dt / 8 * (f - f_n)
    x_c[first_ind] = q_c
    x_c[last_ind] = v_c
    f_c = F(x_c, u_c, params)[last_ind]
    res[last_ind] = v + dt / 6 * (f + 4 * f_c + f_n)
    res[first_ind] = q + dt * v + dt ** 2 / 6 * (f + 2 * f_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
#Numerical values of the paramenters

m0_n, m1_n = [1., 0.3]
l_n = 0.5
g_n = 9.81
params = [g_n, l_n, m0_n, m1_n]
```

```python
#Package imports

import casadi as cas
from optibot.casadi import rhs_to_casadi_function, restriction2casadi
```

```python
# Translating the Sympy Expression into a CasADi function

F_cas_simp = rhs_to_casadi_function(RHS_small[2:], 2)
```

```python
def gen_ini_guess(N = 25, ini_guess = 'lin'):
    '''
    Generates an initial guess for the Cartpole problem of N intervals.
    '''
    if ini_guess == 'zero':
        x_init_guess = np.zeros([N+1,4])
    elif ini_guess == 'lin':
        def_q1 = np.linspace(0,1,N+1)
        def_q2 = np.linspace(0,np.pi,N+1)
        def_v1 = np.zeros(N+1)
        def_v2 = np.zeros(N+1)
        x_init_guess = np.array([def_q1, def_q2, def_v1, def_v2]).T
    return x_init_guess

```

```python
import time
def chrono_solve(opti, solve_repetitions):
    '''
    Calls the solver a certain amount of times and returns the last solution
    obtained and the average computing time
    '''
    cput0 = time.time()
    for ii in range(solve_repetitions):
        sol = opti.solve()
    cput1 = time.time()
    cpudt = (cput1-cput0)/solve_repetitions
    return sol, cpudt

```

```python
#@lru_cache
def casadi_cartpole(N = 25, scheme = 'euler', ini_guess = 'lin', solve_repetitions = 1, t_end = 2):
    opti = cas.Opti()
    p_opts = {"expand":True,'ipopt.print_level':0, 'print_time':0}
    s_opts = {"max_iter": 10000, 'tol': 1e-26}
    opti.solver("ipopt",p_opts,
                        s_opts)
    restr_schemes = {
        'euler': euler_restr, # Euler scheme
        'trapz': trapz_restr, # Trapezoidal Scheme
        'trapz_mod' : trapz_mod_restr, # Second Order Trapezoidal Scheme
        'hs': hs_restr, # Hermite Simpson Scheme, assuming that each Uc is the central value
        'hs_mod': hs_mod_restr, # Second Order Hermite Simpson Scheme, assuming that each Uc is the central value
        'hs_parab': hs_parab_restr, # Hermite Simpson Scheme, with Uc as a free problem parameter
        'hs_mod_parab': hs_mod_parab_restr # Second Order Hermite Simpson Scheme, with Uc as a free problem parameter
        #'your scheme name here': your_scheme_function_here
    }

    f_restr = restr_schemes[scheme]

    # parab is a boolean variable that controls wether the centran points of U are free decision variables
    if scheme in ['hs_parab', 'hs_mod_parab']:
        parab = True
    else:
        parab = False

    # Creating problem structure
    X = opti.variable(N+1,4)
    U = opti.variable(N+1)
    if parab:
        U_c = opti.variable(N)
    T = opti.parameter()
    u_m = opti.parameter()
    Params = opti.parameter(4)

    # Defining the problem cost to minimize (integral of u^2)
    cost = (cas.sum1(U[:]**2)+cas.sum1(U[1:-1]**2))/N
    if parab:
        cost = (4*cas.sum1(U_c[:]**2) + cas.sum1(U[:]**2)+cas.sum1(U[1:-1]**2))/(3*N)
    opti.minimize(cost)

    # Initial and final conditions
    opti.subject_to(X[0,:].T == [0, 0, 0, 0])
    opti.subject_to(X[-1,:].T == [1, np.pi, 0, 0])

    # Translating the scheme restriction function into a CasADi function
    if parab:
        restriction = restriction2casadi(f_restr, F_cas_simp, 2, 1, 4, 1)
    else:
        restriction = restriction2casadi(f_restr, F_cas_simp, 2, 1, 4)

    # Appliying restrictions and action boundaries
    for ii in range(N):
        if parab:
            opti.subject_to(restriction(X[ii,:], X[ii+1,:], U[ii,:], U[ii+1],T/N, Params, U_c[ii])==0)
            opti.subject_to(opti.bounded(-u_m, U_c[ii,:] ,u_m))
        else:
            opti.subject_to(restriction(X[ii,:], X[ii+1,:], U[ii,:], U[ii+1,:],T/N, Params)==0)
        opti.subject_to(opti.bounded(-u_m,U[ii,:],u_m))
    opti.subject_to(opti.bounded(-u_m,U[-1, :],u_m))

    # Setting parameters to their numeric values
    opti.set_value(T, t_end)
    max_f = 20.0
    opti.set_value(u_m, max_f)

    m0_n, m1_n = [1., 0.3]
    l_n = 0.5
    g_n = 9.81
    opti.set_value(Params, [g_n, l_n, m0_n, m1_n])

    # Setting the initialization values
    if ini_guess in ['zero', 'lin']:
        opti.set_initial(X, gen_ini_guess(N, ini_guess))
    elif type(ini_guess) == list:
        opti.set_initial(X, ini_guess[0])
        opti.set_initial(U, ini_guess[1])
        if parab:
            opti.set_initial(U_c, ini_guess[2])
    else:
        raise TypeError('initial guess not understood')

    # Solve
    sol, cpudt = chrono_solve(opti, solve_repetitions)
    err_count = None
    sol_cost = sol.value(cost)
    xx_simp = sol.value(X)
    uu_simp = sol.value(U)
    if parab:
        uu_c = sol.value(U_c)
    else:
        uu_c = None

    # Return data
    return xx_simp, uu_simp, uu_c, cpudt, err_count, sol_cost
```

Let's try to solve the problem for 25 points and the 2nd order Hermite Simpson

```python
from optibot.schemes import interpolated_array, interpolated_array_derivative
from optibot.analysis import dynamic_error
from optibot.numpy import RHS2numpy
```

```python
F_nump = RHS2numpy(RHS_small, 2)
```

```python
scheme = 'hs_mod_parab'
N = 25
xx, uu, uu_c, cpudt, _, cost = casadi_cartpole(N, scheme, 'lin', 1)

xx_interp, uu_interp = interpolated_array(
    X = xx,
    U = uu,
    F = F_nump,
    h = 2/N,
    t_array = np.linspace(0, 2, 2000),
    params = params,
    scheme = "hs_parab",
    u_scheme = 'parab',
    scheme_params = {'u_c' : uu_c}
)
plt.figure(figsize=[16,8])
plt.plot(np.linspace(0,2,N+1),uu[:], 'o',label = '$u_k$ points')
plt.plot(np.linspace(0,2,2*N+1)[1::2],uu_c, 'o',label = '$u_c$ points')
plt.plot(np.linspace(0,2,2000),uu_interp, label = 'interpolation')
plt.grid()
plt.legend()
plt.title('Cart-pole U(t) for 2nd order Hermite Simpson with N = 25')
labels = ['q1','q2','v1','v2']
for ii in range(4):
    plt.figure(figsize=[16,10])
    plt.plot(np.linspace(0,2,N+1),xx[:,ii], 'o',label = f'${labels[ii]}_k$ points')
    plt.plot(np.linspace(0,2,2000),xx_interp[:,ii], label = 'interpolation')
    plt.grid()
    plt.legend()
    plt.title(f'Cart-pole {labels[ii]}(t) for 2nd order Hermite Simpson with N = 25')
```

## Sistematic comparison 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
schemes = ['hs_parab', 'hs_mod_parab', 'trapz', 'trapz_mod'] #If you defined a custom function, name your scheme here
initials = ['lin']
solve_repetitions = 30 #Increase this number to get more reliable values of execution times
N_arr = [20, 25, 30, 40, 50, 60]# You can increase the numbers here, but it will take more time
results = {}

for scheme in schemes:
    for init in initials:
        key = scheme + '_' + init
        print('Problem:', key)
        results[key] = {'N_arr':N_arr}
        for N in N_arr:
            print(f'\tN = {N}')
            xx, uu, uu_c, cpudt, _, cost = casadi_cartpole(N, scheme, init, solve_repetitions)
            results[key][N] = {
                'x': xx,
                'u': uu,
                'u_c': uu_c,
                'cpudt': cpudt,
                'cost': cost,
            }
```

```python
#Calculating the number of collocation number
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']
```

## Dynamic Error

Now we can compute the dynamic errors for each case

```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']
initials = ['lin']#, 'funcs']
n_interp = 4000
for scheme in schemes:
    for init in initials:
        key = scheme + '_' + init
        print('Problem:', key)
        N_arr = results[key]['N_arr']
        for N in N_arr:
            print(f'\tN = {N}')
            if 'parab' in scheme:
                u_scheme = 'parab'
            else:
                u_scheme = 'lin'
            dyn_err_q, dyn_err_v, dyn_err_2_a, dyn_err_2_b = dynamic_error(
                results[key][N]['x'],
                results[key][N]['u'],
                2,
                params,
                F_nump,
                scheme = scheme,
                u_scheme= u_scheme,
                scheme_params={'u_c':results[key][N]['u_c']},
                n_interp = n_interp)
            t_arr = np.linspace(0,2, n_interp)
            tot_dyn_err_q = total_state_error(t_arr, dyn_err_q)
            tot_dyn_err_v = total_state_error(t_arr, dyn_err_v)
            tot_dyn_err_2_a = total_state_error(t_arr, dyn_err_2_a)
            tot_dyn_err_2_b = total_state_error(t_arr, dyn_err_2_b)
            results[key][N]['err_q_int'] = dyn_err_q
            results[key][N]['err_v_int'] = dyn_err_v
            results[key][N]['err_2_a_int'] = dyn_err_2_a
            results[key][N]['err_2_b_int'] = dyn_err_2_b
            results[key][N]['err_q'] = tot_dyn_err_q
            results[key][N]['err_v'] = tot_dyn_err_v
            results[key][N]['err_2_a'] = tot_dyn_err_2_a
            results[key][N]['err_2_b'] = tot_dyn_err_2_b
```

```python
for scheme in schemes:
    for init in initials:
        key = scheme + '_' + init
        print('Problem:', key)
        N_arr = results[key]['N_arr']
        err_q_acum = []
        err_v_acum = []
        err_2_a_acum = []
        err_2_b_acum = []
        cpudt = []
        for N in N_arr:
            err_q_acum.append(results[key][N]['err_q'])
            err_v_acum.append(results[key][N]['err_v'])
            err_2_a_acum.append(results[key][N]['err_2_a'])
            err_2_b_acum.append(results[key][N]['err_2_b'])
            cpudt.append(results[key][N]['cpudt'])
        results[key]['err_q_acum'] = np.array(err_q_acum, dtype = float)
        results[key]['err_v_acum'] = np.array(err_v_acum, dtype = float)
        results[key]['err_2_a_acum'] = np.array(err_2_a_acum, dtype = float)
        results[key]['err_2_b_acum'] = np.array(err_2_b_acum, dtype = float)
        results[key]['cpudt'] = np.array(cpudt, dtype = float)
```

```python
#Plotting parameters
plt.rcParams.update({'font.size': 12})
oct_fig_size = [15,10]
```

```python
sch = [['hs_parab','hs_mod_parab'],['trapz', 'trapz_mod']]
tit = [['Hermite Simpson','2nd order Hermite Simpson'],['Trapezoidal', '2nd order Trapezoidal']]
colors = [f'C{ii}' for ii in [1,0,2,3]]
n_int = len(t_arr)
N_hh = [25,50]
for hh in range(2):
    schemes = sch[hh]
    titles = tit[hh]
    N = N_hh[hh]
    interv_n = (N * t_arr)/2
    for ii in range(2):
        plt.figure(figsize=oct_fig_size)
        for kk in range(len(schemes)):
            scheme = schemes[kk]
            key = scheme + '_lin'
            cut_p = 0
            for ll in range(1,N+1):
                jj = np.searchsorted(interv_n, ll)
                plt.plot(t_arr[cut_p:jj],results[key][N]['err_q_int'][cut_p:jj,ii], '-', c = colors[2*hh+kk], label = titles[kk] if cut_p == 0 else None)
                cut_p = jj
        plt.plot(np.linspace(0,2,N+1), np.zeros(N+1), 'ok', label = 'knot & collocation points')
        if hh == 0:
            plt.plot(np.linspace(0,2,2*N+1)[1::2], np.zeros(N), 'ow', markeredgecolor='k', label = 'collocation points')
        plt.legend()
        plt.grid()
        plt.title(r'First order dynamic error $\varepsilon^{[1]}_{q_'+f'{ii+1}}}$, {titles[0]} schemes, N = {N}')
        plt.xlabel('Time(s)')
        units = 'm/s' if ii == 0 else'rad/s'
        plt.ylabel(f'Dynamic error $({units})$')
        plt.tight_layout(pad = 0.0)
        sch_type = titles[0].replace(' ','_')

        # If you are running the notebook locally and want to save the plots,
        # uncomment the next line
        #plt.savefig(f'Cartpole_First_Order_Dynamic_Error_q_{ii+1}_{sch_type}_schemes_N_{N}.eps', format='eps')
```

```python
sch = [['hs_parab','hs_mod_parab'],['trapz', 'trapz_mod']]
tit = [['Hermite Simpson','2nd order Hermite Simpson'],['Trapezoidal', '2nd order Trapezoidal']]
colors = [f'C{ii}' for ii in [1,0,2,3]]
n_int = len(t_arr)
N_hh = [25,50]
for hh in range(2):
    schemes = sch[hh]
    titles = tit[hh]
    N = N_hh[hh]
    interv_n = (N * t_arr)/2
    for ii in range(2):
        plt.figure(figsize=oct_fig_size)
        for kk in range(len(schemes)):
            scheme = schemes[kk]
            key = scheme + '_lin'
            cut_p = 0
            for ll in range(1,N+1):
                jj = np.searchsorted(interv_n, ll)
                plt.plot(t_arr[cut_p:jj],results[key][N]['err_2_b_int'][cut_p:jj,ii], '-', c = colors[2*hh+kk], label = titles[kk] if cut_p == 0 else None)
                cut_p = jj
        plt.plot(np.linspace(0,2,N+1), np.zeros(N+1), 'ok', label = 'knot & collocation points')
        if hh == 0:
            plt.plot(np.linspace(0,2,2*N+1)[1::2], np.zeros(N), 'ow', markeredgecolor='k', label = 'collocation points')
        plt.legend()
        plt.grid()
        #plt.ylim([-0.00022, 0.00022])
        plt.title(r'Second order dynamic error $\varepsilon^{[2]}_{q_'+f'{ii+1}}}$, {titles[0]} schemes, N = {N}')
        plt.xlabel('Time(s)')
        units = 'm/s^2' if ii == 0 else'rad/s^2'
        plt.ylabel(f'Dynamic error $({units})$')
        plt.tight_layout(pad = 0.0)
        sch_type = titles[0].replace(' ','_')
        # If you are running the notebook locally and want to save the plots,
        # uncomment the next line
        #plt.savefig(f'Cartpole_Second_Order_Dynamic_Error_q_{ii+1}_{sch_type}_schemes_N_{N}.eps', format='eps')
```

```python
schemes_graph = ['hs_mod_parab', 'hs_parab', 'trapz', 'trapz_mod']
titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']
colors = [f'C{ii}' for ii in range(9)]
data_array = ['err_q_acum','err_v_acum','err_2_b_acum','cpudt']
initial = 'lin'

data_key = data_array[2]
for qq in range(2):
    plt.figure(figsize=[10,6])
    plt.title(f'Second order dynamic error $E^{{[2]}}_{{q_{qq+1}}}$')
    for ii in [2,3,1,0]:
        scheme = schemes_graph[ii]
        key = scheme + '_' + initial
        print('Problem:', key)
        N_arr = results[key]['N_arr']
        if len(results[key][data_key].shape) == 1:
            plt.plot(N_arr,results[key][data_key], marker = 'o', c = f'C{ii}',label = titles[ii])
        else:
            plt.plot(N_arr,results[key][data_key][:,qq], marker = 'o', c = f'C{ii}',label = titles[ii])
    plt.yscale('log')
    plt.xlabel('Number of intervals')
    plt.grid()
    plt.legend()
    units = 'm/s' if qq == 0 else'rad/s'
    plt.ylabel(f'Dynamic error $({units})$')
    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'Cartpole_Integrated_Second_Order_Dynamic_Error_q_{qq+1}_vs_N.eps', format='eps')

```

```python
schemes = ['hs_mod_parab','hs_parab', 'trapz', 'trapz_mod']
titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']
plt.figure(figsize=[10,6])
for ii in [2,3,1,0]:
    key = schemes[ii] + '_lin'
    plt.plot(results[key]['N_arr'], results[key][f'cpudt'], marker = 'o', c = f'C{ii}',label = titles[ii])
plt.grid()
plt.legend()
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'Cartpole_optimization_time_vs_interval_number.eps', format='eps')
```

```python
# Here we print the data shown in Table II of the paper
for scheme in ['hs_mod_parab', 'hs_parab', 'trapz', 'trapz_mod']:
    key = scheme + '_lin'
    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]['err_q'], '\n\tErr 2:', results[key][N]['err_2_b'])
```

## Animation

```python
from matplotlib import animation, rc
import matplotlib.patches as patches
from matplotlib.transforms import Affine2D
from IPython.display import HTML
import matplotlib
matplotlib.rcParams['animation.embed_limit'] = 200
```

```python
def create_anim(X, U, params):
    [g_n, l_n, m0_n, m1_n] = params

    N = X.shape[0]
    fig, ax = plt.subplots()
    y_scale = 1
    min_x_cart = np.min(X[:,0])
    max_x_cart = np.max(X[:,0])
    cart_displ = max_x_cart-min_x_cart
    size_x = 2*y_scale + cart_displ
    size_y = 2*y_scale
    draw_width = 14
    draw_height = draw_width / size_x * size_y

    x_0 = X[:,0]
    y_0 = np.zeros_like(x_0)
    x_1 = x_0 + l_n*np.sin(X[:,1])
    y_1 = y_0 - l_n*np.cos(X[:,1])

    x_cm = (m0_n * x_0 + m1_n * x_1)/(m0_n + m1_n)
    y_cm = (m0_n * y_0 + m1_n * y_1)/(m0_n + m1_n)

    fig.set_dpi(72)
    fig.set_size_inches([draw_width,draw_height])
    ax.set_xlim(( min_x_cart-y_scale, max_x_cart+y_scale))
    ax.set_ylim(( -y_scale, y_scale))

    #circle1 = plt.Circle((0, 0), l_n, color='b', ls = ":", fill=False)
    #ax.add_artist(circle1)
    ax.plot([min_x_cart - l_n, max_x_cart + l_n], [0,0], 'k', lw=1, ls = ':')

    line1, = ax.plot([], [], lw=2)
    line3, = ax.plot([], [], 'k', lw=1, ls = ':')
    #line_cm, = ax.plot([], [], 'g', lw=1, ls = ':')
    point0, = ax.plot([], [], marker='s', markersize=10, color="k")
    point1, = ax.plot([], [], marker='o', markersize=7, color="red")
    #point_cm, = ax.plot([], [], marker='o', markersize=10, color="green")
    u_max = max(np.max(np.abs(U[:])),1e-15)
    arrow_w = 0.1*l_n
    arrow_l = 0.7*l_n
    u_arrow = patches.Arrow(0, 0, 0, -arrow_l, color = 'gray',width = arrow_w)
    ax.add_patch(u_arrow)

    print_vars = [X[:,0], X[:,1], U[:], np.linspace(0, N-1, N, dtype=int)]
    print_var_names = ['q_0', 'q_1', 'u_0', 'step']
    texts = []
    ii = 0.8
    for arr in print_vars:
        texts.append(ax.text(-0.8, ii, "", fontsize = 12))
        ii -= 0.2*l_n

    xx_interpolated, uu_interpolated = interpolated_array(
        X,
        U,
        F = F_nump,
        h = 2/(N-1),
        t_array = np.linspace(0, 2, 5*(N-1)+1),
        params = params,
        scheme = 'hs_mod_parab',
        u_scheme = 'parab',
        scheme_params = {'u_c' : results['hs_mod_parab_lin'][N-1]['u_c']}
    )
    x_0_interp = xx_interpolated[:,0]
    y_0_interp = np.zeros_like(x_0_interp)
    x_1_interp = x_0_interp + l_n*np.sin(xx_interpolated[:,1])
    y_1_interp = y_0_interp - l_n*np.cos(xx_interpolated[:,1])

    def init():
        line1.set_data([], [])
        line3.set_data([], [])
        #line_cm.set_data([], [])
        point1.set_data([], [])
        #circle1.center = (0, 0)
        return (line1,)
    def animate(i):
        #circle1.center = (x_0[i], y_0[i])
        point0.set_data(x_0[i], y_0[i])
        line1.set_data([x_0[i], x_1[i]], [y_0[i], y_1[i]])
        point1.set_data(x_1[i], y_1[i])
        #point_cm.set_data(x_cm[i], y_cm[i])
        line3.set_data(x_1_interp[:5*i+1], y_1_interp[:5*i+1])
        #line_cm.set_data(x_cm[:i], y_cm[:i])
        trans = Affine2D()
        u_arrow._patch_transform = trans.scale(U[i] * arrow_l / u_max, arrow_w).translate(x_0[i],0)
        for ii in range(len(texts)):
            text = texts[ii]
            name = print_var_names[ii]
            arr = print_vars[ii]
            if name == 'step':
                text.set_text("$step$ = " + str(arr[i]))
            else:
                text.set_text("$" + name + "$ = %.3f" % arr[i])
        return (line1,u_arrow)
    frame_indices = np.concatenate((np.zeros(10, dtype=int), np.arange(0, N, 1), np.ones(15, dtype=int)*(N-1)))
    anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=frame_indices, interval=20,
                               blit=True)
    return anim
```

```python
anim = create_anim(results['hs_parab_lin'][25]['x'], results['hs_parab_lin'][25]['u'], params)
```

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

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

```python

```

```python

```

```python

```


# Class 5

Note:  The notes that follow are largely those of Mark Krumholz (ANU) who led the Bootcamp
last in 2015.  You can find the 2015 lectures [here](https://sites.google.com/a/ucsc.edu/krumholz/teaching-and-courses/python-15)

```python
# These are to display images in-line
from IPython.display import Image
from IPython.core.display import HTML

#Imports
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
```

Today's class is the final one in the boot camp, and it is here that we turn mainly to scientific applications of python, as opposed to general python programming. The universe of scientific applications is vast, so we will pick out a few examples of common scientific tasks that we can perform using python, numpy, and scipy.

# Root Finding
### Motivating example and statement of the problem

Our first example task is finding the solutions to algebraic equations, or, equivalently, the roots of functions. We can illustrate this example using a problem from basic astronomy. We know that the light emitted by a blackbody follows a distribution in wavelength that depends on its temperature. Specifically, we know that the distribution of intensity versus wavelength is distributed following the Planck function,

\begin{equation}
B_{x}=\left[\frac{2(k T)^{5}}{h^{4} c^{3}}\right] \frac{x^{5}}{e^{x}-1}
\end{equation}

where $\lambda$ is the wavelength, h is Planck's constant, k is Boltzmann's constant, and T is the temperature. Now suppose that we want to find the wavelength of maximum intensity for a given temperature T. This is a straightforward problem of maximizing a function, and we can solve it via the usual calculus method: take the derivative and find where it is equal to zero.

The algebra is a bit less messy if we make the substitutionv $x=hc / \lambda kT$, which makes the function we want to maximize

\begin{equation}
B_{x}=\left[\frac{2(k T)^{5}}{h^{4} c^{3}}\right] \frac{x^{5}}{e^{x}-1}
\end{equation}

Taking the derivative with respect to x, and doing a bit of simplification, we get

\begin{equation}
\frac{d B_{x}}{d x}=\left[\frac{2(k T)^{5}}{h^{4} c^{3}}\right] x^{4} \frac{(5-x) e^{x}-5}{\left(e^{x}-1\right)^{2}}
\end{equation}

We want to find the value of $x$ for which this is equal to 0, which is clearly the value of $x$ where the numerator $(5 - x)e^x - 5$ is equal to zero. We are thus left with a pure math problem: solve the equation

\begin{equation}
(5 - x)e^x - 5 = 0
\end{equation}

The difficulty is that this is a transcendental equation, which means that the equation cannot be solved analytically in an exact way. However, as we will see below, there clearly is a solution.

The problem of solving equations of this goes by the generic name of root-finding. In mathematical terms, the problem can be stated as follows. Suppose we have some function one one variable $f(x)$: for what value(s) of $x$ do we have $f(x) = 0$? In other words, what are the roots of this function?

### Graphical evaluation
A first step toward solving a problem like this is getting some sense of what the function looks like. To that end, we can fire up python, define the function, and plot it. We'll do this for both the Planck function and its derivative, omitting the leading constants in square brackets.

```python
def bx(x):
    return( x**5/(np.exp(x)-1) )
```

```python
def dbdx(x):
    return( x**4 * ((5-x)*np.exp(x)-5) / (np.exp(x)-1)**2 )
```

```python
x = np.arange(0.01,20,0.01)
```

```python
plt.plot(x, bx(x), lw=3)
plt.plot(x, dbdx(x), lw=3)
plt.plot(x, 0*x, 'k--')
plt.legend(['B', 'dB/dx', '0'])
```

The blue line is the graph of the function we want, and the green line is $f(x) = 0$. Clearly there is a solution. By eye, we can estimate that it is around $x = 5$, but we can be much more accurate than that.

One important reason to graph a function before employing these more accurate methods, however, is that a graph can often alert us if a function has more than one root. If it does, we may need to think about which root we want to find, as it is often the case that only one of the roots corresponds to a physically-realistic solution that we're interested in.

## Newton's method and the secant method

So how do we go about finding the numerical value of $x$ for which $f(x) = 0$? There are numerous numerical methods, but in this class we'll explore only two: Newton's method / the secant method, and Brent's method.

Newton's method was invented, as one might guess from the name, by Isaac Newton, and it consists of the following steps. We start with an initial guess for the root $x_0$, and at the guess we evaluate both the function and its derivative: $f(x_0)$ and $f'(x_0)$.

If the function were a straight line, then we could find the root just from this, because $f'(x_0)$ is the slope and $f(x_0)$ is the y value. The root would be at

\begin{equation}
x_1 = x_0 - \frac{f(x_0)}{ f'(x_0)}.
\end{equation}

In general our function f is not a straight line, since, if it were, all of this would be unnecessary. However, on sufficiently small scales, every smooth function looks like a straight line, so this may still be a pretty good approximation, especially if our initial guess was not too far off. To improve the approximation, we can simply repeat the procedure by using the new value $x_1$ to guess a new $x$ value $x_2$ via the same formula:

\begin{equation}
x_2 = x_1 - \frac{f(x_1)}{f'(x_1)}.
\end{equation}

We can perform this operation repeatedly until $f(x_N)$ is as close to zero as we want it to be. The [Newton's method article on wikipedia](https://en.wikipedia.org/wiki/Newton%27s_method)  has a much better graphical representation of this than I am likely to come up with, so here it is:

Newton's method requires that we be able to calculate the derivative of the function $f'(x)$. This is sometimes the case, but now always. If we cannot analytically calculate the derivative, we can use a closely-related method called the second method. This is basically the same as Newton's method, with the difference that the derivative is approximated rather than computed exactly. Suppose that we take $x_0$ and $x_1$ as two initial guesses at the root. We can approximate the derivative by

\begin{equation}
f^{\prime}\left(x_{1}\right) \approx \frac{f\left(x_{1}\right)-f\left(x_{0}\right)}{x_{1}-x_{0}}
\end{equation}

We can plug this approximation into our Newton's method formula to get $x_2$, and then we can approximate the derivative $f'(x_2)$ using the values of $f(x)$ evaluated at $x_2$ and $x_1$ just as we approximated the derivative $f'(x_1)$ using the values of $f(x)$ evaluated at $x_1$ and $x_0$.

The scipy package provides implementations of both Newton's method and the secant method. To use these methods, we must first define the function whose roots we want to find. In the example above, we have already defined the function $f(x)$ that is of interest to us, so we can proceed to the next steps, which are to import the [scipy.optimize](https://docs.scipy.org/doc/scipy/reference/optimize.html) module, and then call the function [newton()](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.newton.html).

The syntax is as follows:

```python
import scipy.optimize as opt

opt.newton(dbdx, 5)
```

    4.965114231744276

The newton() method takes two arguments by default. The first is the function whose root is to be found, and the second is an initial guess for the location of the root. Given this initial guess, the function will apply Newton's method or the secant method iteratively to find the root. In this case, we didn't specify the derivative of the function, and so the secant method is used. If we wanted to specify the derivative, we could do so and then use Newton's method:

```python
def d2bdx2(x):
    return( 20*x**3/(np.exp(x)-1) - 10*x**4*np.exp(x)/(np.exp(x)-1)**2 + 2*np.exp(2*x)*x**5/(np.exp(x)-1)**3 - np.exp(x)*x**5/(np.exp(x)-1)**2 )
```

```python
opt.newton(dbdx, 5, fprime=d2bdx2)
```

    4.965114231744276

The function we have defined as $f'(x)$ is the analytically-computed derivative of $f(x)$. The syntax is that the optional argument fprime is set equal to the function that returns $f'(x)$. Since we have now specified $f'(x)$, the newton() function in scipy.optimize uses Newton's method instead of the secant method. The advantage of Newton's method as opposed to the second method is that it usually converges to the right answer in fewer iterations, so, if the function $f(x)$ takes a while to evaluate, Newton's method can get to the answer faster.

From the output answer, we can find the relationship between temperature and peak wavelength. Recall that we defined

\begin{equation}
x= \frac{h c}{ \lambda k T}
\end{equation}

so we can invert this to give

\begin{equation}
\lambda_{\max } T=x_{\max }(h c / k)=0.290 \mathrm{cm} \mathrm{K}
\end{equation}

where we have plugged in $x = 4.9651$. This is known as **Wien's Law**.

## Bracketed roots, the bisection method, and Brent's method

While Newton's method and the second method worked in the example we just tried, it can also run into problems, and there is no guarantee that it will find a root. Suppose that, instead of using 5 as an initial guess, we had used 0.1. Here's the result:

```python
opt.newton(dbdx, 0.1)
```

So Newton's method failed to find the solution. What went wrong? To answer that question, we need only look at the graph of the function. Newton's method essentially amounts to saying "go downhill / uphill following the current slope", but using 0.1 as an initial guess, "downhill" actually points away from the solution, not toward it. Newton's method works well if we're close enough to the right answer that the function isn't too far from a straight line, but it will fail badly if the function is not like a straight line.

Fortunately, we can do better, if we can bracket the root. Bracketing the root means that we can identify two values a and b such that $f(a)$ and $f(b)$ have different signs meaning that the function $f(x)$ (assuming it is continuous) must pass through 0 for some $x$ between $a$ and $b$.

Bracketing a root is extremely powerful, because it enables one to use an algorithm that is guaranteed to find the root. This method is [bisection](https://en.wikipedia.org/wiki/Bisection_method). The idea of bisection is extremely simple, and can be described by the following steps:

1. Given two points $a$ and $b$ that bracket a root (i.e. $f(a)$ and $f(b)$ have opposite signs), take the halfway point between them $h = (a+b)/2$, and evaluate $f(h)$.

2. If $f(h)$ is within some specified tolerance of 0, then end -- we have found the root.

3. If not, then check if $f(h)$ and $f(a)$ have opposite signs. If they do, then there must be a root between $x = a$ and $x = h$, so set $b = h$ and go back to step 1.

4. If $f(h)$ and $f(a)$ have the same sign, then $f(h)$ and $f(b)$ must have opposite signs, because we know that $f(a)$ and $f(b)$ have opposite signs. Thus there must be a root between $h$ and $b$, so set $a = h$ and go back to step 1.

This method is guaranteed to find the root, because it always keeps the root in the interval of interest, and successively halves that interval until we are close enough.

The bisection method is guaranteed to work, but can be much slower than something like Newton's method, because it only ever gets closer to the answer by a factor of 2 per step. In contrast, if a function is fairly close to a line, Newton's method will jump very close to the answer in a single step. [Brent's method](https://en.wikipedia.org/wiki/Brent's_method) gets the best of both worlds, by trying to use secant or a similar fast method, but falling back on bisection if that doesn't converge toward the answer. It is guaranteed to converge to the right answer just like bisection, but in most cases is nearly as fast as Newton's method.

(Note: although Brent's method is guaranteed to converge *mathematically*, it may still fail to converge *numerically* because on a computer one is using finite-precision arithmetic, and there are situations where the 15 decimal places of accuracy that floating point numbers normally provide will not be enough to get to the answer within the requested tolerance. Attempting to find the root of $f(x) = e^{10000000x} - 1$ numerically is going to be problematic even with an algorithm that is mathematically guaranteed to work.)

The scipy.optimize library implements Brent's method through the routine [brentq()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html), and if should generally be the go-to method of solving equations if you can bracket the root. The syntax is very much like Newton's method, except that, instead of specifying an initial guess, one instead specifies the bracketing values:

```python
opt.brentq(dbdx, 0.1, 100)
```

    4.965114231743836

Here the first argument is the function whose root is to be found, the second argument is the left bracketing value ($a$), and the third argument is the right bracketing value ($b$). Even if these values are nowhere near the root (as is the case here), Brent's method will find it. Calling brentq with values for $a$ and $b$ such that $f(a)$ and $f(b)$ do not have opposite signs results in an error.

## Multi-dimensional root-finding

The problem of finding the roots of a multi-dimensional function, where there is more than one independent or dependent variable, is a significantly harder problem. In its most general form, this is the problem of finding a vector $\mathbf{x}$ satisfying the condition that $\mathbf{f}(\mathbf{x}) = 0$. The most common situation is where $\mathbf{x}$ and $\mathbf{f}(\mathbf{x})$ have the same number of elements, but in general they may have different numbers of elements.

The basic reason that this is a hard problem is that, in more than one dimension, there is no guaranteed-to-work-even-if-it-is-slow method like bisection. There is no way to bracket a root, and guarantee that it lies somewhere between point $a$ and point $b$. The situation is easiest to visualize if we imagine two superimposed landscapes there the elevations of each vary with position. We want to find a point that is at sea level in both landscapes. This is hard. For each landscape we can draw the zero-elevation contour reasonably easily, but there is no guarantee that the two sets of zero-elevation contours ever actually cross one another, and no general way to find out if and where they do.

Due to this difficulty, methods to find the roots of multi-dimensional functions are generally much more complex than those for finding the roots of scalar functions of scalar variables. Indeed, developing algorithms for solving problems of this sort is an active area of research in the field of applied mathematics.

The scipy.optimize module includes several methods for finding the roots of multidimensional functions, most of which can be accessed through the [root()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.html) function. Here's an example of using it. To start with, let us consider a two-dimensional function of two-variables, defined as follows:

```python
def fvec(xvec):
    r = np.sqrt( xvec[0]**2 + xvec[1]**2 )
    phi = np.arccos( xvec[0]/r )
    fx = r**2 - 1.0
    fy = np.cos(phi)
    return( np.array([fx, fy]) )
```

We can see where the solutions are going to be just by looking at the function. Its first component is $r^2 - 1$, so clearly the solution must lie on the circle $r = 1$. The second component is cos(phi), so clearly the solution must lie at phi = $\pi$/2 or $(3/2)\pi$, corresponding to $x = 0$. Thus there are two solutions, one at (0, 1) and one at (0, -1).

To search for these solutions using root, we proceed as follows.

```python
opt.root(fvec, [0.25, 0.5])
```

        fjac: array([[-0.02837482, -0.99959735],
           [ 0.99959735, -0.02837482]])
         fun: array([ 2.69748668e-11, -3.60626471e-11])
     message: 'The solution converged.'
        nfev: 12
         qtf: array([-5.35354922e-09, -4.22230880e-09])
           r: array([-0.96382325,  0.02222859,  2.06183641])
      status: 1
     success: True
           x: array([-3.60625931e-11,  1.00000000e+00])

The first argument to root is the function whose roots are to be found, and the second argument is an initial guess at the solution. There are several keywords that control what method is used to search for the solution, and things like that. The function returns an object with a number of parts describing whether it succeeded in finding a solution, and if so where. In this example, success is True, indicating that the solver did find a solution, and x is the solution it found. In this case, it found the solution (0, 1).

# Numerical Integration

### Motivating example and statement of the problem

The second topic for today is numerical integration: using a numerical method to approximately evaluate integrals. To motivate this problem, we can return to the Planck function and the distribution of energy from a radiating black body, and ask an interesting question. Human eyes are sensitive to light over the wavelength range (roughly) 400 - 700 nm. What fraction of the Sun's light output actually occurs over this range of wavelengths. Now the Sun isn't really a blackbody, but its spectrum is close enough to being a blackbody that we can get a good rough answer by treating it as one.

The Sun's effective temperature is 5780 K, so using our definition $x = h c / \lambda k T$, plugging in the values of [Planck's constant](http://en.wikipedia.org/wiki/Planck_constant), [Boltzmann's constant](http://en.wikipedia.org/wiki/Boltzmann_constant), and the [speed of light](http://en.wikipedia.org/wiki/Speed_of_light), we find that

\begin{equation}
x=\frac{24.89}{(\lambda / 100 \mathrm{nm})}
\end{equation}

Thus the range 400 - 700 nm corresponds to $x = 3.56 - 6.22$. The distribution of energies follows the Planck function, so we are interested in knowing how the integral of the Planck function over this interval compares to the integral from 0 to infinity. In mathematical terms, we want to be able to evaluate

\begin{equation}
\frac{\int_{3.56}^{6.22} B_{x} d x}{\int_{0}^{\infty} B_{x} d x}
\end{equation}

he bottom integral can in fact be done exactly analytically (the result is $16 \pi^{6}(k T)^{5} /\left(63 h^{4} c^{3}\right)$ ), but the top integral cannot be evaluated analytically, and must instead be evaluated numerically. We can represent the problem geometrically with a simple plot:

```python
x1 = np.arange(3.56, 6.22, 0.001)

plt.clf()
plt.fill_between(x, bx(x), alpha=0.5)
plt.fill_between(x1, bx(x1), alpha=0.5, facecolor='r')
plt.ylim([0,25])
```

The goal is to evaluate the ratio of the purple shaded area to the total blue plus purple shaded area.

This is an example of a general problem that can be stated very simply: given a known function $f(x)$, evaluate the definite integral

\begin{equation}
\int_{a}^{b} f(x) d x
\end{equation}

over some specified interval (a,b), where a could be -infinity, and b could be +infinity.

### Quadrature of specified functions

There are numerous methods for evaluating integrals, and many of them are implemented in the [scipy.integrate](http://docs.scipy.org/doc/scipy/reference/integrate.html) module. All of these use some variant on the idea of quadrature, which means to approximate the function using a series of other functions whose integrals are easy to evaluate. A trivial example of this is the trapezoidal rule, where one approximates the function by a series of straight lines. In these case, the shape formed by the line segment on top, the y axis on the bottom, and the vertical sides is a trapezoid, a shape whose area is trivial to calculate. The integral may therefore be approximated as the sum of the areas of a series of trapezoids.

Modern quadrature methods are much more clever than this, in that they tend to approximate each segment with somewhat higher order functions than simple lines, and that they adaptively decide where to put the breaks between the various approximate segments based on an estimate of the numerical error. The scipy routine [quad()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad.html) implements one such sophisticated algorithm. Its use is quite simple. To evaluate the numerator of our integral, we can do

```python
import scipy.integrate as integ

integ.quad(bx, 3.56, 6.22)
```

    (53.14174673016606, 5.899919078862899e-13)

The syntax is of quad is that the first argument is the function to be integrated, the second argument is the lower limit of integration, and the third argument is the upper limit. The function returns a tuple of two numbers. The first is the value of the integral, and the second is an estimate of the absolute value of the error in the result.

The limits of integration can be plus or minus infinity, which are represented by Inf and -Inf in python. Thus to evaluate the denominator, we could do

```python
integ.quad(bx, 0, np.Inf)
```

    /home/bruno/pyEnvs/py3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp

    (122.08116743813392, 2.0046050699847875e-07)

It is easy to verify that this matches the analytic result very precisely, with an error that is in fact much smaller than the estimate given by the quad routine:

```python
integ.quad(bx, 0, np.Inf)[0]-8*np.pi**6/63
```

    /home/bruno/pyEnvs/py3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp

    4.263256414560601e-14

Thus the answer to our question, "what fraction of the Sun's light falls in the wavelength range visible to the human eye?", is

```python
integ.quad(bx, 3.56, 6.22)[0] / integ.quad(bx, 0, np.Inf)[0]
```

    /home/bruno/pyEnvs/py3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: overflow encountered in exp

    0.43529848088237094

Thus about 44% of the Sun's output falls in the wavelength range that the human eye can see. At first this might seem surprising, given how tiny a part of the spectrum we can see. Of course, however, this is no accident -- evolution would hardly select for vision in parts of the spectrum where there isn't much light to see. The eye have evolved to be exquisitely well-tuned to the light output by the Sun. We would do much less well around stars with different effective temperatures.

### Quadrature of sampled functions

In the example we just did, we had an analytic formula for the function we want to integrate. However, that is not always the case. Sometimes we only have access to the value of the function evaluated at particular points. For example, the function values might be data that were obtained experimentally rather than by an exact calculation, or they might be the result of a complex numerical calculation that we can't afford to perform a very large number of times. Obviously in this case we are not going to get as good an approximation to an integral as if we knew the function perfectly and could evaluate it wherever we wanted. However, we would still like to be able to numerically integrate such sampled functions as well as possible.

As an example, we can again take on the question of what fraction of the Sun's output energy lies in the energy range from 400 - 700 nm, but this time using a measured Solar spectrum rather than a pure blackbody. To start with, download the file [sorce_ssi.csv](https://sites.google.com/a/ucsc.edu/krumholz/teaching-and-courses/python14/class-5/sorce_ssi.csv?attredirects=0&d=1), which contains a measured spectrum of the Sun from January 1, 2010, obtained from the very useful database maintained by the [SOlar Radiation and Climate Experiment (SORCE)](http://lasp.colorado.edu/home/sorce/). This file contains three columns. The first is observation time (all the same in this file), the second is wavelength in nm, and the third is radiation flux in W/m2/nm.

Let's start by reading in this data and plotting it. We will do so using the astropy.io package:

```python
from astropy.io import ascii

data = ascii.read('sorce_ssi.csv')
data.columns
```

    <TableColumns names=('time (days since 2003-01-24)','wavelength (nm)','irradiance (W/m^2/nm)')>

The columns available are as we said: time, wavelength, and irradiance, another word for frequency-dependent flux. We can plot the latter two quantities against one another as follows:

```python
l = data['wavelength (nm)']
flux = data['irradiance (W/m^2/nm)']
plt.plot(l, flux)
```

is is similar to our idealized blackbody, but not quite identical. To find out the fraction of the power between 400 and 700 nm, we want to integrate this function over that range, and divide by the integral of the full range.

Python provides a method to evaluate these tabulate integrals by approximating the function over each tabulation range as a parabola. This approximation method is known as Simpson's rule, and the routine that performs it is [scipy.integrate.simps()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.simps.html). We can use it as follows:

```python
integ.simps(flux, x=l)
```

    1319.4359600971845

The first argument is the $y$ values of the function to be integrated, and the optional argument $x =$ lets us specify the $x$ values.

This gives the integral over the full range of the function. To get the integral over 400 - 700 nm, we need to extract just the parts of the l and flux arrays that correspond to this. In order to do this, we need to find out what indices in x correspond to values of 400 and 700. There are many ways to do this in numpy, but the easiest is probably the following:

```python
l1 = l[np.logical_and(l>400, l<700)]
flux1 = flux[np.logical_and(l>400, l<700)]
```

Let's unpack these statements a bit. The statement l > 400 produces an array that is True wherever l is bigger than 400, and False otherwise. Similarly, l < 700 produces an array that is True wherever l is below 700, False elsewhere. Then we take logical_and(l>400, l<700), and the effect of logical_and is exactly what it sounds like: it produces True where both conditions (l>400 and l<700) are satisfied, and False elsewhere. Finally, l[logical_and(l>400, l<700)] finds the elements of l where both conditions are satisfied. Thus we've set l1 to just the elements of l that are between 400 and 700. Similarly, the second statement sets flux1 to be equal to just the elements of flux where the corresponding value of l is between 400 and 700. We can verify that this all worked ok by checking the maximum and minimum values of l1:

```python
np.amax(l1)
```

    698.85

```python
np.amin(l1)
```

    400.34

Finally, we are ready to use simps() to evaluate the integral of flux1 versus l1, and compare that to the integral of flux versus l:

```python
integ.simps(flux1, x=l1) / integ.simps(flux, x=l)
```

    0.40204419522708396

So we get 40% instead of 44% using the real Solar spectrum -- not a big difference, clearly.

### Multi-dimensional quadrature

The quadrature method we have used thus far can also be generalized to multiple dimensional integrals, although of course the cost of the evaluation goes up as the dimensionality does. This capability is provided by the functions dblquad (for 2d functions), tplquad (for 3d functions), and nquad (for arbitrary-dimensional quadratures); see scipy.integrate. We won't go over the usage of these functions in class, but they're basically analogous to what we've already done.

# Ordinary Differential Equations

### Motivating example and statement of the problem

Our final example in this class is using python to integrate ordinary differential equations. The example we will use is a simple harmonic oscillator, for example a mass on a spring. Let us consider a mass m attached to a spring with spring constant $k$. The spring is oriented horizontally, so there is no gravitational force. When the mass is displaced by a distance $x$ from the position where the spring is force-free, the restoring force exerted on the mass is $-kx$. This in turn causes the mass to accelerate with an acceleration $a = -(k/m) x$. We are interested in figuring out how the mass will move if we initially displace it some distance $x_0$ from the equilibrium position.

The first step is to write down the equations governing this system. Let $v = dx/dt$ be the velocity of the mass at any given time. The set of equations that we have to solve is then

\begin{aligned}
&\frac{d x}{d t}=v\\
&\frac{d v}{d t}=-\frac{k}{m} x
\end{aligned}

Each of these equations is an ordinary differential equation (ODE), and this pair represents a coupled pair of ODEs. Our goal is to find functions $x(t)$ and $v(t)$ that satisfy these two equations, along with the initial condition that $x(0) = x_0$ and $v(0) = 0$.

This particular set of ODEs can be solved analytically, but many others can't be. In general such a system can be written as follows. A single ODE is an equation of the form

\begin{equation}
\frac{d x}{d t}=f(x, t)
\end{equation}

where $f(x, t)$ is some arbitrary, specified function of $x$ and $t$. A set of two coupled ODEs is a pair of equations of the form

\begin{aligned}
&\frac{d x}{d t}=f(x, y, t)\\
&\frac{d y}{d t}=g(x, y, t)
\end{aligned}

where $f(x,y,t)$ and $g(x,y,t)$ are any specified functions. ODEs of this form are very common in physics, and usually instead of $y$ we have the velocity, $v$. In general a system of N ODEs takes the form

\begin{aligned}
&\frac{d x_{1}}{d t}=f_{1}\left(x_{1}, x_{2}, x_{3}, \dots, t\right)\\
&\frac{d x_{2}}{d t}=f_{2}\left(x_{1}, x_{2}, x_{3}, \dots, t\right)\\
&\frac{d x_{3}}{d t}=f_{3}\left(x_{1}, x_{2}, x_{3}, ., t\right)\\
&...
\end{aligned}

where there are N independent variables, and N specified functions $f_1$ through $f_N$. This system of equations requires a set of $N$ initial conditions, which are most commonly of the form $x_i(0) = x_{i,0}$.

## Solving ODEs in python

There are an immense number of standard methods for solving ODEs, and scipy provides an interface to a powerful and general ODE solver. As with integration, the first step is to define the function specifying the equations to be integrated. Specifically, we need to provide a function that gives the right-hand sides of the above equations, that is the derivatives of $x$ and $v$ with respect to time. We can enter this as

```python
def derivs(xv, t, k, m):
    x = xv[0]
    v = xv[1]
    dxdt = v
    dvdt = -(k/m)*x
    return( [dxdt, dvdt] )
```

The form of this function has to be as follows. The first argument is an array giving all the dependent variables to be integrated in time. In this case, our two dependent variables are $x$ and $v$, so we get an array of two elements, the first one being $x$ and the second one being $v$. The second argument of derivs is the time $t$. We don't actually need the time in this case, but we get it anyway because in general the derivatives might depend on time. Then finally we have any additional arguments we need. In this case the additional arguments are the spring constant $k$ and the mass $m$.

The function must then return an array giving the derivatives of the two dependent variables. Thus the function has to return $dxdt$ and $dvdt$. The function we have written calculates these from the equations we are trying to solve.

Once we've entered the function, the next thing to do is to specify the initial conditions. In this example, since there are two dependent variables, $x$ and $v$, we need to give the initial values of $x$ and $v$. As we described the problem, the initial position is displaced from $x = 0$ by some amount, which for this example we can take to be 1, and the initial velocity is zero. Thus we say

```python
xv0 = [1, 0]
```

The third step is to decide at what times we'd like to know the position and velocity. Thus is up to us, but for this example let's say we'd like to know the position at intervals of 0.1 seconds for 20 seconds. We will create an array specifying these times:

```python
t = np.arange(0, 20, 0.1)

```

Finally, we need to choose values for the mass and spring constant. Let's choose a mass of 1 kg, and a spring constant of 2 N/m. Thus

```python
k = 2.
m = 1.
```

Now we're finally ready to integrate. We use the function [odeint()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html) as follows:

```python
output = integ.odeint(derivs, xv0, t, args=(k, m))
```

The first argument is the function specifying the derivatives, the second is the initial conditions, the third is the set of times at which we want to know the position and velocity. Finally, the optional argument args specifies the additional arguments (in this case $k$ and $m$) that are to be passed to the derivs function.

The odeint routine returns at output that consists of the position $x$ and velocity $v$ at every time we asked for. The shape is

```python
output.shape
```

    (200, 2)

Thus output is a 200 x 2 array. To plot the position versus time and velocity versus time, we can do

```python
plt.clf()
plt.plot(t, output[:,0])
plt.plot(t, output[:,1])
plt.legend(['x', 'v'])
```

Thus the position and velocity both undergo sinusoidal oscillations, $\pi/2$ out of phase, as we expect for this simple harmonic oscillator.


# 15.5. A bit of number theory with SymPy

https://ipython-books.github.io/155-a-bit-of-number-theory-with-sympy/

## Ref

* Undergraduate level: Elementary Number Theory, Gareth A. Jones, Josephine M. Jones, Springer, (1998)
* Graduate level: A Classical Introduction to Modern Number Theory, Kenneth Ireland, Michael Rosen, Springer, (1982)
* SymPy's number-theory module, available at http://docs.sympy.org/latest/modules/ntheory.html
* The Chinese Remainder Theorem on Wikipedia, at https://en.wikipedia.org/wiki/Chinese_remainder_theorem
* Applications of the Chinese Remainder Theorem, given at http://mathoverflow.net/questions/10014/applications-of-the-chinese-remainder-theorem
* Number theory lectures on Awesome Math, at https://github.com/rossant/awesome-math/#number-theory

```python
from sympy import *
import sympy.ntheory as nt
init_printing()
```

```python
nt.isprime(11), nt.isprime(121)
```

    (True, False)

```python
nt.isprime(2017)
```

    True

```python
nt.nextprime(2017)
```

```python
nt.prime(1000)
```

```python
nt.primepi(2017)
```

```python
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
x = np.arange(2, 10000)
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
ax.plot(x, list(map(nt.primepi, x)), '-k',
        label='$\pi(x)$')
ax.plot(x, x / np.log(x), '--k',
        label='$x/\log(x)$')
ax.legend(loc=2)
```

```python
factors = nt.factorint(2020)
```

```python
type(factors), factors.keys(), factors.values()
```

    (dict, dict_keys([2, 5, 101]), dict_values([2, 1, 1]))

```python
n = 1
for k,v in factors.items():
    n *= k**v
print(n)
```

    2020

```python
nt.factorint(1998)
```

```python
2 * 3**3 * 37
```

```python
from sympy.ntheory.modular import solve_congruence
solve_congruence((1, 3), (2, 4), (3, 5))
```

```python

```


# Maximum likelihood Estimation (MLE)
based on http://python-for-signal-processing.blogspot.com/2012/10/maximum-likelihood-estimation-maximum.html
## Simulate coin flipping
- [Bernoulli distribution](https://en.wikipedia.org/wiki/Bernoulli_distribution)
is the probability distribution of a random variable which takes the value 1 with probability $p$ and the value 0 with probability $q = 1 - p$
- [scipy.stats.bernoulli](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.bernoulli.html)

```python
import numpy as np
from scipy.stats import bernoulli

np.random.seed(123456789)

p_true = 1/2 # this is the value we will try to estimate from the observed data
fp = bernoulli(p_true)

def sample(n=10):
    """
    simulate coin flipping
    """
    return fp.rvs(n) # flip it n times

xs = sample(100) # generate some samples
```

## Find maximum of Bernoulli distribution
Single experiment
$$\phi(x) = p ^ {x} * (1 - p) ^ { 1 - x }$$
Series of experiments
$$\mathcal{L}(p|x) = \prod_{i=1}^{n} p^{x_{i}}*(p-1)^{1-x_{i}}$$
### Hints
- [sympy.diff()](http://docs.sympy.org/dev/modules/core.html#sympy.core.function.diff)
- [sympy.expand()](http://docs.sympy.org/dev/modules/core.html#sympy.core.function.expand)
- [sympy.expand_log()](http://docs.sympy.org/dev/modules/core.html#sympy.core.function.expand_log)
- [sympy.solve()](http://docs.sympy.org/dev/modules/core.html#sympy.core.function.solve)
- [sympy.symbols()](http://docs.sympy.org/dev/modules/core.html#symbols)
- [sympy gotchas](http://docs.sympy.org/dev/tutorial/gotchas.html)

```python
import sympy
from sympy.abc import x

p = sympy.symbols('p', positive=True)
phi = p ** x * (1 - p) ** (1 - x)
L = np.prod([phi.subs(x, i) for i in xs]) # objective function to maximize
log_L = sympy.expand_log(sympy.log(L))
sol = sympy.solve(sympy.diff(log_L, p), p)[0]
```

```python
import matplotlib.pyplot as plt

x_space = np.linspace(1/100, 1, 100, endpoint=False)

plt.plot(x_space,
         list(map(sympy.lambdify(p, log_L, 'numpy'), x_space)),
         sol,
         log_L.subs(p, sol),
         'o',
         p_true,
         log_L.subs(p, p_true),
         's',
        )
plt.xlabel('$p$', fontsize=18)
plt.ylabel('Likelihood', fontsize=18)
plt.title('Estimate not equal to true value', fontsize=18)
plt.grid(True)
plt.show()
```

## Empirically examine the behavior of the maximum likelihood estimator
- [evalf()](http://docs.sympy.org/dev/modules/core.html#module-sympy.core.evalf)

```python
def estimator_gen(niter=10, ns=100):
    """
    generate data to estimate distribution of maximum likelihood estimator'
    """
    x = sympy.symbols('x', real=True)
    phi = p**x*(1-p)**(1-x)
    for i in range(niter):
        xs = sample(ns) # generate some samples from the experiment
        L = np.prod([phi.subs(x,i) for i in xs]) # objective function to maximize
        log_L = sympy.expand_log(sympy.log(L))
        sol = sympy.solve(sympy.diff(log_L, p), p)[0]
        yield float(sol.evalf())

entries = list(estimator_gen(100)) # this may take awhile, depending on how much data you want to generate
plt.hist(entries) # histogram of maximum likelihood estimator
plt.title('$\mu={:3.3f},\sigma={:3.3f}$'.format(np.mean(entries), np.std(entries)), fontsize=18)
plt.show()
```

## Dynamic of MLE by length sample sequence

```python
def estimator_dynamics(ns_space, num_tries = 20):
    for ns in ns_space:
        estimations = list(estimator_gen(num_tries, ns))
        yield np.mean(estimations), np.std(estimations)

ns_space = list(range(10, 100, 5))
entries = list(estimator_dynamics(ns_space))
entries_mean = list(map(lambda e: e[0], entries))
entries_std = list(map(lambda e: e[1], entries))

plt.errorbar(ns_space, entries_mean, entries_std, fmt='-o')
plt.show()
```

```python

```


# <ins>Lab #2: Basics + Loops, Conditionals & Functions </ins>

These exercises are meant to give you some practice with the concepts from Tutorials 2.1-2.3, as well as Tutorials 1.1/1.2 since we didn't get to the exercises last week. Be sure to use best practices, comment your code, and be mindful of your units.

```python
# Import any libraries here:
import math
import numpy as np
import matplotlib.pyplot as plt
```

***
## Arithmetic & variables

1. Use the quadratic equation, to find the roots of $5x^2 + 3x + 1$ by calculating the numerator and denominator separately. Round to 2 decimal places. (**Tip:** When parintheses in a big equation get a little hard to read, breaking it up like this is a good way to make code easier to read.)

$$x = \frac{-b \pm \sqrt{b^2 - 4ac}}{2a}$$

```python
# Your code here:

import math
a,b,c = 5,3,-1

num_plus = -b + math.sqrt(b**2 - 4 * a * c)
num_minus = -b - math.sqrt(b**2 - 4 * a * c)
denom = 2*a

root_plus = num_plus/denom
root_minus = num_minus/denom

print(round(root_plus,2), round(root_minus,2))
```

    0.24 -0.84

2. Use the distance modulus, to find the distance to a star with apparent magnitude $m_v = 0.5$ and absolute magnitude $M_v = -5.85$.

$$m_v - M_v = 5log(d/10 \text{ pc})$$

```python
# Your code here:

m, M = 0.5, -5.85
d = 10**((m-M)/5)*10

print(d, 'pc')
```

    186.20871366628677 pc

***
## User input & data types

Prompt the user for their age and calculate the percent of their life they've completed assuming they live to 80 years old.

```python
# Your code here:

age = input('How old are you in years? ')

print((int(age)/80.0)*100.0, '%')
```

    How old are you in years? 21
    26.25 %

***
## Indexing

1. Print the strings "pneumono," "ultra," "volcano," and "coniosis" by slicing the word "pneumonoultramicroscopicsilicovolcanoconiosis".

```python
# Your code here:

string = "pneumonoultramicroscopicsilicovolcanoconiosis"

fourth,ninth,fifteenth = string[3], string[9], string[15]
index_20, index_3, index_12 = string[20], string[3], string[12]

print("Fourth letter:", fourth)
print("Ninth letter:", ninth)
print("Fifteenth letter:", fifteenth)

print("Third index:", index_3)
print("Twelfth index:", index_12)
print("Twentieth index:", index_20)
```

    Fourth letter: u
    Ninth letter: l
    Fifteenth letter: c
    Third index: u
    Twelfth index: a
    Twentieth index: o

2. Print the 4th, 9th, and 15th letter and the letter at the 20th, 3rd, and 12th indices of the word from part one by indexing.

```python
# Your code here:

string = "pneumonoultramicroscopicsilicovolcanoconiosis"

pneumono = string[:8]
ultra = string[8:13]
microscopic = string[13:24]
volcanoconiosis = string[-15:]

print(pneumono, ultra, microscopic, volcanoconiosis)
```

    pneumono ultra microscopic volcanoconiosis

***
## `numpy` arrays

Use the Wein displacement law to calculate the temperatures of blackbodies with 20 peak wavelengths between 300nm and 700nm using arrays.

\begin{equation}
    \lambda_{peak}T = 0.29 \text{ cm K}
\end{equation}

```python
# Your code here:

lambdas = np.linspace(300,700,20)
lambdas_cm = lambdas*1E-7

temps = 0.29 / lambdas_cm

print(temps)
```

    [9666.66666667 9032.78688525 8476.92307692 7985.50724638 7547.94520548
     7155.84415584 6802.4691358  6482.35294118 6191.01123596 5924.7311828
     5680.41237113 5455.44554455 5247.61904762 5055.04587156 4876.10619469
     4709.4017094  4553.71900826 4408.         4271.31782946 4142.85714286]

***
## Conditional statements

Write a program that has the user input a wavelength or frequency and returns (a) the photon energy, and (b) the part of the electromagnetic spectrum in which a photon of that energy falls.

$$E = h\nu \quad\text{&}\quad \lambda\nu = c$$

```python
# Your code here:

wav_or_freq = input("Would you like to enter a wavelength or a frequency? Enter w for wavelength, f for frequency: ")
val = float(input("Enter the value in centimeters for wavelength, in Hertz for frequency: "))

if wav_or_freq == 'w':
    val = val #cm
    energy = 4.13E-15*(3*10**10)/(val) #eV
elif wav_or_freq == 'f':
    energy = 4.13E-15*val #eV
    val = (3*10**10)/(val) #cm

if val < 1E-9:
    print("Gamma; energy = ", energy, "eV")
elif val < 1.0E-6 and val > 1.0E-9:
    print("X-ray; energy = ", energy, "eV")
elif val < 3.5E-5 and val > 1.0E-6:
    print("UV; energy = ", energy, "eV")
elif val < 7.5E-5 and val > 3.5E-5:
    print("Visible; energy = ", energy, "eV")
elif val < 0.5 and val > 7.5E-5:
    print("IR; energy = ", energy, "eV")
else:
    print("Radio; energy = ", energy, "eV")

```

    Would you like to enter a wavelength or a frequency? Enter w for wavelength, f for frequency: w
    Enter the value in centimeters for wavelength, in Hertz for frequency: 1
    Radio; energy =  0.0001239 eV

***
## Loops and plotting

1. Add together every third number in the list [155,2,54,34,5,16,7,38,26,10] and print the result.

```python
# Your code here:

lst = [155,2,54,34,5,16,7,38,26,10]

summation = 0
for n,x in enumerate(lst):
    if n % 3 == 0:
        summation += x

print(summation)
```

    206

2. Plot the sine and cosine functions on separate sides of the y-axis.

```python
# Your code here:

x = np.linspace(-2*np.pi,2*np.pi)
sine = np.sin(x)
cosine = np.cos(x)

plt.plot(x,sine)
plt.plot(x,cosine)
```

3. The rules for Pin-Pon are as follows:

    - Start counting numbers, starting from 1.
    - Whenever the next number's a multiple of 3 (3, 6, 9,…), replace the actual number with, "Pin."
    - Whenever the next number's a multiple of 5, replace it with the word, "Pon."
    - Whenever the number's a multiple of both 3 and 5, say, "Pin Pon."
    - Keep going until someone messes up.

Write a program that prints the results of the game for 1 through 100.

```python
# Your code here:

for n in range(1,101):
    if (n % 3 == 0) and (n % 5 == 0):
        print("Pin Pon")
    elif n % 5 == 0:
        print("Pin")
    elif n % 3 == 0:
        print("Pon")
    else:
        print(n)
```

***
## Putting it all together

These are meant to use all of your new skills at once. That being said, you may have notice that there's no "Functions" section above --- once you've completed each of the tasks below, please turn them into functions, comment them well, and add them to a Python script which will serve as a library for the future.

1. (From Lab 1): While DMS (degrees, minutes, seconds) and HMS (hours, minutes, seconds) formats for celestial coordinates have their place (e.g. in every database ever), having them in decimal degrees is often more convenient for calculations. Write two scripts: one that allows the user to enter in a DMS coordinate and prints it in decimal degrees, and another for HMS coordinates.

```python
# Your code here:

# DMS
dms = input("Enter D:M:S coordinates, separated by colons: ")
dms = dms.split(":")
degs, mins, secs = dms
decdeg = float(degs) + float(mins)/60.0 + float(secs)/3600.0

print("Decimal degrees:", decdeg)

# HMS
hms = input("Enter H:M:S coordinates, separated by colons: ")
hms = hms.split(":")
hrs, mins, secs = hms
decdeg = (float(hrs) + float(mins)/60.0 + float(secs)/3600.0) * 15

print("Decimal degrees:", decdeg)
```

2. In my research on X-ray binaries, I utilize the Novikov-Thorne thin accretion disk approximation for black holes, which relates the temperature of the accertion disk to the mass of the accretor. Typical disk temperatures range between 0.5 and 1.0 keV. Using this relationship, calculate the mass of black holes (a) within this range, (b) in the "interesting" regime of 1.0 to 5.0 keV, and (c) in the "crazy" range of >5.0 keV. Plot your results, making sure the three regimes are clearly labeled in your plot. (P.S. -- X-ray astronomers often talk about temperatures in terms of energy rather than Kelvin, so keep that in mind with this equation.)

$$T_{eff} \text{ [keV]} \sim \Bigg(\frac{10M_\odot}{M}\Bigg)^{1/4} \text{ keV}$$

```python
# Your code here:

temps_normal = np.linspace(0.5,1.0,100)
temps_interesting = np.linspace(1.0,5.0,100)
temps_crazy = np.linspace(5.0,6.0,100)

'''
Input: temp -- keV
Output: mass -- solar masses
'''
def mass(temp):
    return (10.0/temp**4)**(-1.0)

mass_normal = mass(temps_normal)
mass_interesting = mass(temps_interesting)
mass_crazy = mass(temps_crazy)

plt.plot(temps_normal,mass_normal,label="Normal")
plt.plot(temps_interesting,mass_interesting,label="Interesting")
plt.plot(temps_crazy,mass_crazy,label="Crazy")
plt.xlabel("Disk temperature [keV]")
plt.ylabel(r"Mass of black hole $M_\odot$")
plt.title("Novikov-Thorne Thin Accretion Disk Relationship")
plt.legend()
plt.show()
plt.close()
```

```python

```


# PCA

```python
import pandas
# For lots of great things.
import numpy as np
# To make our plots.
import matplotlib.pyplot as plt
%matplotlib inline
# Because sympy and LaTeX make
# everything look wonderful!
from sympy import *
init_printing(use_latex=True)
from IPython.display import display
# We will use this to check our implementation...
from sklearn.decomposition import PCA
import keras
from sklearn.preprocessing import StandardScaler
```

```python
## load in data and split into x train and y train
## data = np.array(pandas.read_csv("./comp_new_trainingdata.csv", header=0))
data = np.array(pandas.read_csv("./training_noavg.csv", header=0))
## bring in loc 0 -1 test data for PCA
data1 = np.array(pandas.read_csv("./test1.csv", header=0))
data2 = np.array(pandas.read_csv("./test2.csv", header=0))
data3 = np.array(pandas.read_csv("./test3.csv", header=0))
## Have to drop all teh rows that have nan values because they will not help with net
## clean out rows with nan values
data = data[~np.isnan(data).any(axis=1)]
data1 = data1[~np.isnan(data1).any(axis=1)]
data2 = data2[~np.isnan(data2).any(axis=1)]
data3 = data3[~np.isnan(data3).any(axis=1)]

print(data[:8])
print(data.shape)

data = np.vstack((data,data1,data2,data3))
print(data[:8])
data.shape
```

```python
# vectors AND class labels...
X = data[:,0:8] # 0 thru 30
Y = data[:,8] # 30

scaler = StandardScaler()
# standardize X .. will mean center data
X = scaler.fit_transform(X)

# Pretty-print with display()!
display(X.shape)
display(Y.shape)
display(Matrix(np.unique(Y)).T)
display(X[0:8])
```

```python
U,S,V = np.linalg.svd(X,full_matrices=True)

# Percent variance accounted for
plt.plot(100.0*S/np.sum(S))
plt.ylabel('% Var')
plt.xlabel('Singular Value')
plt.show()
```

```python
# Variance accounted for in the first two principal components
100.0*(S[0]+S[1])/np.sum(S)
```

```python
# Scale the singular vectors, resulting in a rotated form of our mean-centered data
D = np.zeros([X.shape[0],X.shape[1]])
np.fill_diagonal(D,S)
Xrotated = np.dot(U,D)
# Extract just the first two principal components!
PCs = Xrotated[:,0:2]
PCs.shape
```

```python
# The x and y values come from the two
# Principal Components and the colors for
# each point are selected based on the
# corresponding class for each point...
plt.scatter(PCs[:,0],PCs[:,1],
color=[['red','green','blue','orange'][i] for i in Y.astype(int)])
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()
```

The data suggest that we have some clear descision boundries.  The orange represents the testing dat from all three locations, we can see that it fits the class groupings.  A simple MLP with linear activation functions will not work for our data, we should use a RELU or sigmoid instead along with a large hidden layer.  It remains to be seen how well the network will be able to generalize.


# Basics of Hafnians and Loop Hafnians
*Author: Nicolás Quesada*

In the [background section](../hafnian.html) of the The Walrus documentation, some basic ideas related to (loop) hafnians were introduced. This tutorial is a computational exploration of the same ideas.

```python
from thewalrus.reference import hafnian as haf_ref
from thewalrus import hafnian
import numpy as np
import matplotlib.pyplot as plt
%config InlineBackend.figure_formats=['svg']
```

## A simple loopless graph and the hafnian

Let's consider the following graph

with adjacency matrix

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

It is easy to verify by inspection that the graph in Fig. 1 has only one perfect matching given by the edges (1,4)(2,5)(3,6).
We can verify this by calculating the hafnian of the adjacency matrix $A$

```python
haf_ref(A) # Using the reference implementation
```

    1

```python
hafnian(A) # Using the default recursive method
```

    1

Let's see what happens if we rescale the adjacency matrix by a scalar $a$. We'll use the [SymPy](https://sympy.org) library for symbolic manipulation:

```python
from sympy import symbols
```

```python
a = symbols("a")
haf_ref(a*A)
```

    a**3

The example above shows that one can use the reference implementation not only with numpy arrays but also with symbolic sympy expressions.

## A graph with loops and the loop hafnian

Now let's consider a graph with loops:

The adjacency matrix is now

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

Note that now the adjacency matrix has non zero elements in the diagonal.
It is also strightforward to see that the graph in Fig. 2 has two perfect matchings, namely, (1,4)(2,5)(3,6) and (1,1)(5,5)(2,4)(3,6)

```python
haf_ref(At, loop=True) # Using the reference implementation
```

    2

```python
hafnian(At, loop=True) # Using the default recursive method
```

    2.0000000000000107

We can also use the loop hafnian to count the number of matching (perfect or otherwise)
by taking the adjacency matrix of the loop less graph, putting ones on its diagonal and calculating the loop hafnian of the resulting matrix. For the graph in Fig. 1 we find

```python
haf_ref(A+np.diag([1,1,1,1,1,1]), loop=True)
```

    15


```python
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
import Euler as p1
import nonlinear_ODE as p2
```

# Exercise C.2 Solving a nonlinear ODE
Solves the ODE problem: $$u^\prime = u^q,\ u(0) = 1,\ t\epsilon [0,T]$$

$q=1,\ \Delta t=0.1:$

```python
p2.graph(1, 0.1, 1)
```

$q=1,\ \Delta t=0.01:$

```python
p2.graph(1, 0.01, 1)
```

$q=2,\ \Delta t=0.1:$

```python
p2.graph(2, 0.1, 1)
```

$q=12,\ \Delta t=0.01:$

```python
p2.graph(2, 0.01, 1)
```

$q=3,\ \Delta t=0.1:$

```python
p2.graph(3, 0.1, 1)
```

```python
p2.graph(3, 0.01, 1)
```

The Euler Method is most accurate for small $\Delta t$, which corresponds to a larger number of calculations, and small $q$.


# Generating the cluster number density

We will compute the cluster number density, $n(s, p)$, from sets of percolating two-dimensional clusters. As the cluster number density is not known analytically in all but the simplest systems, i.e., the infinite dimensional and the one-dimensional system, we will estimate it numerically. This is done by
\begin{align}
    n(s, p; L) \approx \frac{N_s}{ML^d},
\end{align}
where $s$ is the size of the cluster, $p$ the probability for a site to be set in the system, $L$ the length of a side in the system (assuming equal lengths for all sides), $d$ the dimension of the problem ($L^d$ is the volume), $N_s$ the number of clusters of size $s$ and $M$ the number of simulations performed.

```python
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.ndimage
import skimage
import tqdm
import sklearn.linear_model
from uncertainties import ufloat

sns.set(color_codes=True)
```

```python
def is_percolating(prop, num_rows, num_cols):
    min_row, min_col, max_row, max_col = prop.bbox

    return max_row - min_row == num_rows or max_col - min_col == num_cols
```

```python
def compute_cluster_area(system, p, remove_percolating_cluster=True):
    mat = system < p
    # Label and count the number of connected clusters
    labels, num_features = scipy.ndimage.measurements.label(mat)
    s_list = skimage.measure.regionprops(labels)

    new_s_list = []
    if remove_percolating_cluster:
        for s in s_list:
            if is_percolating(s, *system.shape):
                continue
            new_s_list.append(s)
    else:
        new_s_list = s_list

    area = list(map(lambda prop: prop.area, new_s_list))

    return area
```

Example code to plot a specific prop in a system.

```python
box = s[1].bbox
plt.matshow(labels[box[0]:box[2], box[1]:box[3]], cmap="hsv")
plt.show()
```

```python
def compute_cluster_number_density(L, M, p, a=1.2):
    area = []
    for i in range(M):
        z = np.random.rand(L, L)
        area.extend(compute_cluster_area(z, p))

    n, s = np.histogram(area, L ** 2)

    # Logarithmic binning
    logamax = np.ceil(np.log(max(s)) / np.log(a))
    bins = a ** np.arange(0, logamax, 1)

    nl, _ = np.histogram(area, bins)
    ds = np.diff(bins)
    sl = (bins[1:] + bins[:-1]) * 0.5
    nsl = nl / (M * L ** 2 * ds)

    # Remove bins where there are no observations
    mask = np.abs(nsl) > 1e-14

    return sl[mask], nsl[mask]
```

```python
p_c = 0.59275
```

Below we plot the cluster number density as $p \to p_c = 0.59275$ from below.

```python
fig = plt.figure(figsize=(14, 10))

for p in tqdm.tqdm_notebook(np.linspace(0.3, p_c, 6)):
    plt.loglog(
        *compute_cluster_number_density(200, 200, p, a=1.2),
        label=rf"$p = {p:.5f}$"
    )

plt.legend(loc="best")
plt.title(r"Cluster number density for $p \to p_c$ from below")
plt.xlabel(r"$s$")
plt.ylabel(r"$n(s, p)$")
plt.show()
```

In this figure we can see how the cluster number density follows a seemingly straight line for a while, before dropping of sharply when $p$ is a little off from the critical percolation probability. This shows how the cluster number density at the critical percolation probability follows a power law as the axes are logarithmic.

Here we repeat the experiment, but for $p \to p_c$ from above.

```python
fig = plt.figure(figsize=(14, 10))

for p in tqdm.tqdm_notebook(np.linspace(p_c, 0.7, 6)):
    plt.loglog(
        *compute_cluster_number_density(200, 200, p, a=1.2),
        label=rf"$p = {p:.5f}$"
    )

plt.legend()
plt.title(r"Cluster number density for $p \to p_c$ from above")
plt.xlabel(r"$s$")
plt.ylabel(r"$n(s, p)$")
plt.show()
```

When $p \to p_c$ from above, we see that all cluster number densities behave in a power law fashion. Why is that?

## The cluster number density at the critical percolation probability

For $p = p_c = 0.59275$, we explore how the cluster number density changes for a system of $L = 2^k$ for $k = \{4, \dots, 9\}$.
We wish to show how the cluster number density deviates from the power law behavior, when $p = p_c$. We expect that
\begin{align}
    n(s, p_c; L) = F\left(\frac{s}{s_{\xi}}\right) s^{-\tau} \propto s^{-\tau},
\end{align}
where $F$ is some function that is constant when $p = p_c$.
By taking the logarithm on both sides, we get an expression for $\tau$.
\begin{gather}
    \log\left[n(s, p_c; L)\right]
    =
    -\tau \log\left[ s \right]
    + \log\left[ F\left(\frac{s}{s_{\xi}}\right) \right].
\end{gather}
We use linear regression to get an estimate for $-\tau$. We also compute $-\tau$ manually by computing the decrease between two values of $s$, for all differences. We find our best estimate by computing the mean of $-\tau$.

```python
L_arr = 2 ** np.arange(4, 10)
true_tau = 187 / 91
```

```python
fig = plt.figure(figsize=(14, 10))
a = 1.2
tau_list = []

for L in tqdm.tqdm_notebook(L_arr):
    sl, nsl = compute_cluster_number_density(L, 500, p=p_c, a=a)

    log_sl = np.log(sl) / np.log(a)
    log_nsl = np.log(nsl) / np.log(a)

    plt.loglog(sl, nsl, label=rf"$L = {L}$")
    clf = sklearn.linear_model.LinearRegression(
        fit_intercept=True
    ).fit(log_sl[:, None], log_nsl[:, None])

    tau = clf.coef_[0, 0]
    C = clf.intercept_[0]

    print(f"For L = {L}: -tau = {tau:.4f}\tC = {C:.4f}")

    tau_arr = np.zeros(int(len(log_nsl) * (len(log_nsl) - 1) / 2))
    index = 0
    for i in range(len(log_nsl)):
        for j in range(i + 1, len(log_nsl)):
            tau = (log_nsl[j] - log_nsl[i]) / (log_sl[j] - log_sl[i])
            tau_arr[index] = tau
            index += 1

    tau_list.append(tau_arr.copy())

plt.title(r"Cluster number density as a function of $s$ in order to estimate $\tau$")
plt.legend()
plt.xlabel(r"$s$")
plt.ylabel(r"$n(s, p)$")
plt.show()
```

Below we display the mean value for $-\tau$ with its standard deviation for each size $L$.

```python
for L, tau in zip(L_arr, tau_list):
    print(f"For L = {L}: tau = {np.mean(tau):.2f} +/- {np.std(tau):.2f}")
```

    For L = 16: tau = -1.83 +/- 0.57
    For L = 32: tau = -1.93 +/- 0.68
    For L = 64: tau = -1.89 +/- 0.58
    For L = 128: tau = -1.86 +/- 0.33
    For L = 256: tau = -1.88 +/- 0.28
    For L = 512: tau = -1.90 +/- 0.27

We see that this is quite far off from the true value $-\tau = 187 / 91 \approx 2.05$. Below we output the error in $-\tau$.

```python
unc_tau = [ufloat(np.mean(tau), np.std(tau)) for tau in tau_list]

for L, tau in zip(L_arr, unc_tau):
    print(f"For L = {L}: error tau = {true_tau + tau}")
```

    For L = 16: error tau = 0.2+/-0.6
    For L = 32: error tau = 0.1+/-0.7
    For L = 64: error tau = 0.2+/-0.6
    For L = 128: error tau = 0.20+/-0.33
    For L = 256: error tau = 0.18+/-0.28
    For L = 512: error tau = 0.15+/-0.27

We can see that our estimate for $-\tau$ is quite far off. This indicates that using linear regression on the slope of the logarithmically binned cluster number density against the cluster size is not very precise.

## Measuring the characteristic cluster size $s_\xi$

The scaling of the characteristic cluster size $s_{\xi}$ is given by
\begin{align}
    s_{\xi} \propto |p - p_c|^{-1 / \sigma}.
\end{align}
We _define_ the characteristic cluster size by the point where
\begin{gather}
    \frac{n(s, p)}{n(s, p_c)} = F\left(\frac{s}{s_{\xi}}\right) = \frac{1}{2},
    \\
    \implies
    n(s, p) = \frac{1}{2} n(s, p_c).
\end{gather}
The equality is hard to match exactly, so numerically we define it to be
\begin{align}
    n(s, p) \leq \frac{1}{2} n(s, p_c).
\end{align}
The value of $s$ that first satisfies this equation is then our estimate of the characteristic cluster size $s_{\xi}$. We look at the characteristic cluster size for $p \to p_c$ from below.

```python
L = 512
M = 500
p_list = [0.45, 0.5, 0.54, 0.57, 0.58] # Numbers from book
```

```python
fig = plt.figure(figsize=(14, 10))
a = 1.3
s_xi_list = []
sl_list = []
nsl_list = []

sl_c, nsl_c = compute_cluster_number_density(L, M, p=p_c, a=a)
plt.loglog(sl_c, nsl_c, label=rf"$p = {p_c}$")

for p in tqdm.tqdm_notebook(p_list):
    sl, nsl = compute_cluster_number_density(L, M, p=p, a=a)
    sl_list.append(sl)
    nsl_list.append(nsl)

    plt.loglog(sl, nsl, label=rf"$p = {p}$")

    mask = nsl <= 0.5 * nsl_c[:len(nsl)]

    # If no points are below the cluster number density at p = p_c,
    # skip that index.
    if not np.any(mask):
        continue

    index = np.argmax(mask)

    s_xi = sl[index]
    s_xi_list.append(s_xi)

    plt.scatter(s_xi, nsl[index], label=(r"$s_\xi = {0:.2}$").format(s_xi))

plt.legend()
plt.xlabel(r"$s$")
plt.ylabel(r"$n(s, p)$")
plt.title(r"Estimating the characteristic cluster size")
plt.show()
```

In this plot we have marked the characteristic cluster size $s_{\xi}$ by dots. We now plot these as a function of the percolation probability $p$.

```python
fig = plt.figure(figsize=(14, 10))

plt.plot(p_list, s_xi_list, "-o")
plt.xlim(0.4, 0.7)
plt.title(r"The characteristic cluster size $s_{\xi}$ as function of percolation probability $p$")
plt.xlabel(r"$p$")
plt.ylabel(r"$s_{\xi}$")
plt.show()
```

In this figure we see how the characteristic cluster size $s_{\xi}$ diverges as $p \to p_c$ from below, that is, the cluster number density follows the power law type behaviour for longer before dropping off towards zero. For the scaling
\begin{align}
    s_{\xi} = C|p - p_c|^{-1 / \sigma},
\end{align}
where $C$ is a constant, we can find an expression for the exponent $\sigma$ by taking the logarithm on both sides and rearranging the terms.
\begin{gather}
    \log(s_{\xi}) = \log(C) - \frac{1}{\sigma}\log\left(|p - p_c|\right).
\end{gather}
We use linear regression to find an expression for the coefficient $-\sigma^{-1}$.

```python
s_xi_arr = np.log(np.array(s_xi_list))
p_min_pc = np.log(np.abs(p_c - np.array(p_list)))
```

```python
clf = sklearn.linear_model.LinearRegression().fit(
    p_min_pc[:, None], s_xi_arr[:, None]
)

coef = clf.coef_[0, 0]
sigma = - 1 / coef

print(f"sigma = {sigma}")
```

    sigma = 0.47692140021677076

The true value for $\sigma$ in $d = 2$ dimensions is
\begin{align}
    \sigma = \frac{36}{91} \approx 0.395.
\end{align}
The absolute difference between this value and the one found by us is computed below.

```python
sigma_true = 36 / 91

print(f"Diff in sigma: {np.abs(sigma_true - sigma)}")
```

    Diff in sigma: 0.08131700461237518

```python
fig = plt.figure(figsize=(14, 10))

for i, p in enumerate(p_list):
    y_vals = nsl_list[i] * sl_list[i] ** true_tau
    x_vals = sl_list[i] * np.abs(p - p_c) ** (1 / sigma)

    plt.loglog(x_vals, y_vals, label=fr"$p={p}$")

plt.legend(loc="best")
plt.title(r"Data collapse for $n(s, p)$")
plt.xlabel(r"$s |p - p_c|^{1/\sigma}$")
plt.ylabel(r"$n(s, p) s^{\tau}$")
plt.show()
```

In this figure we can see the data collapse of the cluster number density.

## Mass scaling of percolation cluster

We now compute the mass of the percolating cluster $M(L)$ at $p = p_c$.
When $p = p_c$ the characteristic length, $\xi \to \infty$.
This means that the mass scales as
\begin{align}
    M(L) \propto L^{D}.
\end{align}
Taking the logarithm on both sides, we can find an estimate for $D$ using linear regression, viz.
\begin{align}
    \log\left(M(L)\right)
    = \log(C) + D\log(L),
\end{align}
where $C$ is a constant.
We compute the mass of the percolating cluster by counting the number of set sites contained in the cluster that spans the entire system length $L$.

```python
def compute_mass_of_percolating_cluster(system):
    L = system.shape[0]

    # Label and count the number of connected clusters
    labels, num_features = scipy.ndimage.measurements.label(system)
    s_list = skimage.measure.regionprops(labels)

    for s in s_list:
        if is_percolating(s, *system.shape):
            # We look at the first percolating cluster that we find.
            # There might be more than one, but we ignore these.
            return s.area

    return 0
```

```python
L_arr = 2 ** np.arange(4, 11)
mass_arr = np.zeros(len(L_arr))
M = 100
```

We repeat the experiment $M$ times to make sure that we actually get a percolating cluster for all system sizes.

```python
for i, L in tqdm.tqdm_notebook(enumerate(L_arr)):
    for _ in range(M):
        system = np.random.rand(L, L) <= p_c
        mass_arr[i] += compute_mass_of_percolating_cluster(system)

mass_arr /= M
```

    Widget Javascript not detected.  It may not be installed or enabled properly.

Fitting the data to find $D$ and the intercept $C$.

```python
clf = sklearn.linear_model.LinearRegression().fit(
    np.log(L_arr)[:, None],
    np.log(mass_arr)[:, None]
)

D = clf.coef_[0, 0]
print(f"D = {D}")
```

    D = 1.895808803079029

The true value of $D$ is given by
\begin{align}
    D = \frac{91}{48} \approx 1.896.
\end{align}
Thus, the absolute difference between the true value and our estimated value is given by.

```python
D_true = 91 / 48

print(f"Diff: {np.abs(D_true - D)}")
```

    Diff: 2.453025430426692e-05

```python
fig = plt.figure(figsize=(14, 10))

plt.loglog(L_arr, mass_arr, "-o", label=r"Data")
plt.loglog(
    L_arr,
    np.exp(clf.intercept_[0]) * L_arr ** D,
    "--",
    label=r"$M(L) = C L^{D}$"
)

plt.title(r"Mass scaling of the percolating cluster as a function of side length $L$")
plt.xlabel(r"$L$")
plt.ylabel(r"$M(L)$")
plt.legend()
plt.show()
```

In this figure we can see how the log-log plot of the mass and the system size follows a power law type behaviour. We have also plotted the theoretical line on top using our estimate for $D$.