# Closed loop simulations

In this notebook we solve the control problem at the vehicle level. The algorithm is a centralized predictive controller. The notebook is divided in multiple sections

- [Vehicle dynamics](#dynamics)
- [Control problem](#control_problem)
    - [Optimality conditions](#optimal_conditions)
- [Functions at operational level](#compute_control)
- [Connection with symuvia and data storage](#closed_loop)

## System dynamics

In addition the dynamics of the system are defined as

\begin{equation}
\begin{bmatrix}
\dot{x}_i\\
\dot{s}_i\\
\dot{v}_i\\
\end{bmatrix}=
\begin{bmatrix}
v_i\\
v_{i-1}-v_i\\
u_i
\end{bmatrix}
\end{equation}

## Control problem

In this case the solution to the optimal control problem corresponds to the solution of the optimal problem:

\begin{equation}
\begin{aligned}
    \underset{\mathbf u}{\min}\ J\big(\mathbf{x},\mathbf{u},t|\mathbf{x}(0)\big) =& \underset{\mathbf u}{\min}& &\int^{T_p}_0 \mathcal{L}(\mathbf x (\tau), \mathbf u(\tau),\tau)\quad d\tau + \mathcal{G}(\mathbf x(T_p),T_p)\\
    & \text{s.t} & & \mathbf x(0)= \mathbf x_0\\
    & & & \dot{\mathbf x} = \mathbf f(\mathbf x, \mathbf u)
\end{aligned}
\end{equation}

for a determined cost function $\mathcal{L}(\mathbf x, \mathbf u)$.

## Cost function

In this case the cost function is defined as follows:

\begin{equation}
\begin{aligned}
\mathcal{L}(s_i,v_i) = C_1 \big(s_i-v_i\cdot t_h+s_0\big)^2 + C_2\big(v_i-v_{i-1}\big)^2+C_3u^2_i
\end{aligned}
\end{equation}

where $s_i$ represents the spacing of $i$-vehicle, $v_i$ its speed and $a_i$ the vehicle acceleration, $u_i$ is the external acceleration to be applied by the control. The parameters $t_h, s_0$ represent the time headway, and constant space reference.

#### Optimality conditions

The solution of the optimal control problem requires the implementation of the Hamiltonian given by
\begin{equation}
    \mathcal{H}(\mathbf x,\mathbf u) = \mathcal{L}(\mathbf x,\mathbf u) +\lambda^T\mathbf{f}(\mathbf x,\mathbf u)
\end{equation}

For the given boundary value problem (BVP), two simultaneous initial value problems (IVP) are solved simultaneously. The solution that satisfies
\begin{equation}
    \mathbf u^\star = \min \mathcal{H}(\mathbf x, \mathbf u, \lambda), \quad \text{s.t}\ \mathbf u \in \mathcal{U}, x \in \mathcal{X}
\end{equation}
when the second differential equation is satisfied by the by the co-state dynamic equation
\begin{equation}
\begin{aligned}
-\frac{d}{dt}\lambda =& \frac{\partial\mathcal H}{\partial \mathbf x} = \frac{\partial \mathcal L}{\partial \mathcal x}+\lambda \cdot \frac{\partial \mathbf f}{\mathbf x}\\
\lambda(T_p)&=\frac{\partial \mathcal G}{\partial x}(\mathbf x(T_p))
\end{aligned}
\end{equation}

#### Optimality conditions (2nd order dynamics)

The hamiltonian can be written this time as:

\begin{equation}
\begin{aligned}
\mathcal{H}(s_i,v_i,u_i) =& C_1 \left(s_i-(v_i\cdot t_g+s_0)\right)^2 + C_2\big(v_{i-1}-v_{i}\big)^2+C_3u^2_i+ \lambda_s(v_{i-1}-v_{i})+\lambda_v u_i
\end{aligned}
\end{equation}

The conditions lead to the following conditions in the case of the following conditions:

\begin{equation}
    \begin{aligned}
    -\frac{\partial \lambda_s}{\partial t} & = \frac{\partial \mathcal H}{\partial s_i} = 2C_1\left(s_i -(v_it_g+s_0)\right) & \rightarrow &\dot{\lambda}_s = -2C_1\left(s_i -(v_it_g+s_0)\right) \\
    -\frac{\partial \lambda_v}{\partial t} & = \frac{\partial  \mathcal H}{\partial v_i} = -2C_1\left(s_i -(v_it_g+s_0)\right)t_g-2C_2(v_{i-1}-v_{i})-\lambda_s & \rightarrow &\dot{\lambda}_v = 2C_1\left(s_i -(v_it_g+s_0)\right)t_g + 2C_2(v_{i-1}-v_{i})+\lambda_s
    \end{aligned}
\end{equation}

In addition,

\begin{equation}
\begin{aligned}
    u^\star_i &= \frac{\partial \mathcal H}{\partial u_i} = 2C_3u_i + \lambda_v = 0\\
    u^\star_i &= -\frac{\lambda_v}{2c_3}
\end{aligned}
\end{equation}

Package importing

```python
import os

from sqlalchemy import create_engine, MetaData
from sqlalchemy import Table, Column, String, Integer, Float
from sqlalchemy import insert, delete, select, case, and_

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Bokeh
from bokeh.plotting import figure, show
from bokeh.sampledata.iris import flowers
from bokeh.io import output_notebook
from bokeh.palettes import Viridis
output_notebook()

import ipywidgets as widgets
from IPython.display import display
```

    Loading BokehJS ...

#### Creation of tables

In this case table `headway` and `traj` should be already filled according to notebooks [01-Open-loop.ipynb](01-Open-loop.ipynb) and [02-Tactical-strategy.ipynb](02-Tactical-strategy.ipynb)

```python
dir_path = os.getcwd()

engine_path = ('..','Output','SymOut.sqlite')
engine_name = os.path.join(os.path.sep,*engine_path)
engine_full_name = os.path.join(dir_path,*engine_path)
engine_call = 'sqlite://'+engine_name
engine = create_engine(engine_call)
metadata = MetaData()

if os.path.isfile(engine_full_name):
    try:
        ltbstr = 'Loaded table in: '
        connection = engine.connect()
        traj = Table('traj', metadata, autoload=True, autoload_with=engine)
        headway = Table('headway', metadata, autoload=True, autoload_with=engine)
        closed = Table('closed', metadata, autoload=True, autoload_with=engine)
        control = Table('control', metadata, autoload=True, autoload_with=engine)
        stmt = delete(closed)
        results = connection.execute(stmt)
        stmt = delete(control)
        results = connection.execute(stmt)
    except:
        ltbstr = 'Created table in: '
        closed = Table('closed', metadata,
                 Column('ti', Float()),
                 Column('id', Integer()),
                 Column('type', String(3)),
                 Column('tron', String(10)),
                 Column('voie', Integer()),
                 Column('dst', Float()),
                 Column('abs', Float()),
                 Column('vit', Float()),
                 Column('ldr', Integer()),
                 Column('spc', Float()),
                 Column('vld', Float()))
        control= Table('control', metadata,
                 Column('ti', Float()),
                 Column('id', Integer()),
                 Column('type', String(3)),
                 Column('tron', String(10)),
                 Column('voie', Integer()),
                 Column('ctr', Float()),
                 Column('nit', Integer()))
        metadata.create_all(engine)
        connection = engine.connect()
    finally:
        print(ltbstr, engine)

```

    Loaded table in:  Engine(sqlite:///../Output/SymOut.sqlite)

#### Functions for operational level

The following auxiliary functions are created in order to solve the problem in the function [compute_control](#compute_control)

- `reversedEnumerate` Consider a list of `np.array` and flips the array along the 0 axis
- `initial_setup_mpc` Initializes variables for the controller
- `find_idx_ldr` Finds the leader of vehicles at time t.
- `format_reference` converts the reference information form a db into a matrix.

```python
def reversedEnumerate(*args):
    """ Inverse enumeration iterator"""
    revArg = [np.flip(x, axis = 0) for x in args]
    return zip(range(len(args[0])-1, -1, -1), *revArg)

def find_idx_ldr(results):
#     """ From dbQuery finds idx or leader for CAVs"""
    # Frozen network (A-priori)
    ldrl = [dveh_ldr[x[1]] for x in results if x[2]=='CAV']
    idx_ldr = [dveh_idx[x] for x in ldrl]

    return idx_ldr, ldrl

def initial_setup_mpc(results, h_ref):
    """ Initialize variables for controller
    """

    TGref = format_reference(h_ref)
    h = TGref.shape[0]

    n_CAV = len([ty[2] for ty in results if ty[2]=='CAV'])
    dCAVu = [h, n_CAV]
    # print(f'Dimensions control: {dCAVu}')

    Sref = np.zeros(dCAVu)

    S = np.zeros(dCAVu)
    V = np.zeros(dCAVu)
    DV = np.zeros(dCAVu)
    Lv = np.zeros(dCAVu)
    Ls = np.zeros(dCAVu)
    return (Sref, TGref, S, V, DV, Ls, Lv)

def format_reference(h_ref):
    """ Convert query from a reference into a
        numpy array
    """
    vehids = set([x[1] for x in h_ref])

    # Rearrange
    refDf = pd.DataFrame(h_ref, columns = ['ti','id','gapt'])
    # Pivot to pass vehicles as columns
    refMat = pd.pivot_table(refDf, index='ti', columns='id')['gapt']
    refMat = refMat.as_matrix()

    return refMat
```

#### Control definition

The function computes the optimal control for a set of vehicles

```python

def compute_control(results, h_ref, u_lead):

    Sref, Tgref, S, V, DV, Ls, Lv = initial_setup_mpc(results, h_ref)

    # Static leadership
    ldr_pos, ldr_list = find_idx_ldr(results)

    S0 = [s[9] for s in results if s[2]=='CAV']
    V0 = [v[7] for v in results if v[2]=='CAV']
    DV0 = [dv[10]-dv[7] for dv in results if dv[2]=='CAV']
    U_ext = Lv
    #U_ext[:,0] = u_lead # Head acceleration (external)

    # Initialize global variables
    S[0] = S0
    V[0] = V0
    DV[0] = DV0
    h = len(S)
    n = 0
    n_prev = 0

    # Parameters
    C1 = 0.1
    C2 = 1
    C3 = 0.5
    ALPHA = 0.01
    1
    EPS  = 0.1
    error = 100

    bSuccess = 2
    N = 100001 # number of iterations
    step = iter(range(N))

    while (error > EPS) and (bSuccess>0):
        try:
            next(step)
            U_star = -Lv/(2*C3)
            U_star = np.clip(U_star, U_MIN, U_MAX)

            DU = U_star[:,ldr_pos]-U_star[:] + U_ext

            # Forward evolution
            for i,u_s,du in zip(range(h), U_star, DU):
                if i<len(S)-1:
                    DV[i+1] = DV[i]+ DT * du
                    S[i+1] = S[i] + DT * DV[i]
                    V[i+1] = V[i]+ DT * u_s

            Sref = V * Tgref + 1/KC

            # Forward plots
            # plot_forward(Sref, Tgref, S, V, DV, U_star)

            ls = np.zeros(Ls.shape)
            lv = np.zeros(Lv.shape)

            # Backward evolution
            for i, s, v, dv, tg in reversedEnumerate(S, V, DV, Tgref):
                if i>0:
                    sref = v * tg + 1/KC
                    lv[i-1] = lv[i] + DT * (-2 * C1 * (s-sref) * tg - C2 * dv - ls[i])
                    ls[i-1] = ls[i] + DT * (2 * C1 * (s-sref))

            # Update
            Ls = (1 - ALPHA) * Ls + ALPHA * ls
            Lv = (1 - ALPHA) * Lv + ALPHA * lv

            # Backwards plots
            # plot_backwards(ls, lv, Ls, Lv)

            error = np.linalg.norm(Ls - ls) + np.linalg.norm(Lv-lv)
            # print(f'Iteration: {n}, Error: {error}')

            # Routine for changing convergence parameter

            if error > 10e5:
                raise AssertionError('Algorithm does not converge ')
            if n >= 500:
                ALPHA = max(ALPHA - 0.01, 0.01)
                #print(f'Reaching {n} iterations: Reducing alpha: {ALPHA}')
                #print(f'Error before update {error}')
                if n > 10000:
                    raise AssertionError(
                        'Maximum iterations reached by the algorithm')
                n_prev = n + n_prev
                n = 0
            if error <= EPS:
                bSuccess = 0

            n += 1

        except StopIteration:
            print('Stop by iteration')
            print('Last simulation step at iteration: {}'.format(n+n_prev))
            bSuccess = 0

    n = n + n_prev
    return (S, V, DV, U_star, DU, n)

def determine_lane_change(CAVabsP):
    """ Returns the tuple (tron, voie) for a
        CAV vehicle based on positions updates.
    """

    # Works only in current network

    CAVtron = []
    CAVvoie = []

    for abs_x in CAVabsP:
        if (abs_x <= 0):
            CAVtron.append('In_main')
            CAVvoie.append(1)
        elif (abs_x > 0) and (abs_x <=100.0):
            CAVtron.append('Merge_zone')
            CAVvoie.append(2)
        else:
            CAVtron.append('Out_main')
            CAVvoie.append(1)
    return CAVtron, CAVvoie

def update_state(S, V, DV, U_star, DU, n, results_closed):
    """ Updates the state and computes closed loop updates
    """

    # NOTE: To be taken into account. Closed loop simulations
    # run without Symuvia. Requires implementation of the connection
    # NO LANE CHANGE MODEL IMPLENTED FOR HDV

     # Forward evolution
    DVp = DV[0] + DT * DU[0]
    Sp = S[0] + DT * DV[0]
    Vp = V[0] + DT * U_star[0]

    # t_i, id, type, from (results_closed):
    # Keep the same along the simulation
    CAVti = [x[0]+ DT for x in results_closed if x[2]=='CAV']
    CAVti = [float(np.round(x,1)) for x in CAVti]
    CAVid = [x[1] for x in results_closed if x[2]=='CAV']
    CAVtype = [x[2] for x in results_closed if x[2]=='CAV']

    # Updates

    # Postition query
    CAVdst = [x[5] for x in results_closed if x[2]=='CAV']
    CAVabs = [x[6] for x in results_closed if x[2]=='CAV']

    # Updates from closed loop
    CAVdstP = [x+ DT * v for (x,v) in zip(CAVdst, V[0])] #or Vp?
    CAVabsP = [x+ DT * v for (x,v) in zip(CAVabs, V[0])] #or Vp?

    # Lane change
    CAVtron, CAVvoie = determine_lane_change(CAVabsP)

    # State updates
    ldr_pos, ldr_list = find_idx_ldr(results_closed)
    CAVvitP = Vp
    CAVldr = ldr_list
    CAVldr = [int(x) for x in ldr_list]
    CAVspcP = Sp
    CAVvldP = DVp + Vp

    # Render CAV results for dB
    CAViter = zip(CAVti, CAVid, CAVtype, CAVtron, CAVvoie,
                  CAVdstP, CAVabsP, CAVvitP, CAVldr,
                  CAVspcP, CAVvldP)

    keys = ('ti', 'id', 'type', 'tron', 'voie',
            'dst', 'abs', 'vit', 'ldr',
            'spc','vld')

    keysU = ('ti', 'id', 'type', 'tron', 'voie',
             'ctr', 'nit')

    lVehTrajCL = []
    for i in CAViter:
        lVehTrajCL.append(dict(zip(keys,i)))

    n_list = [n]*U_star.shape[0]

    CAViter = zip(CAVti, CAVid, CAVtype, CAVtron,
                  CAVvoie, U_star[0], n_list)

    lVehUCL = []
    for i in CAViter:
        lVehUCL.append(dict(zip(keysU,i)))

    return lVehTrajCL, lVehUCL

def format_open_loop(results):
    """ Aux function
        To write in the closed loop database
        results from open loop without treatment
        (No control applied)
        Homogenizes results in terms of content
    """

    keys = ('ti', 'id', 'type', 'tron', 'voie',
            'dst', 'abs', 'vit', 'ldr',
            'spc','vld')
    keysU = ('ti', 'id', 'type', 'tron', 'voie',
             'ctr', 'nit')

    lVehTrajOL = []
    lVehUOL = []
    for i in results:
        lVehTrajOL.append(dict(zip(keys,i)))
        ti, vid, ty, tr, vo, _, _, _, _, _, _ = i
        u_tup = (ti, vid, ty, tr, vo, 0, 0)
        lVehUOL.append(dict(zip(keysU,u_tup)))

    return lVehTrajOL,lVehUOL

```

### Launch closed loop simulation

Closed loop simulations

```python
progressSim = widgets.FloatProgress(
    value=5,
    min=0,
    max=71.9,
    step=0.1,
    description='Simulating:',
    bar_style='info',
    orientation='horizontal'
)
```

Parameters

```python
DT = 0.1 # Sample time

KC = 0.16 # CAV max density
KH = 0.0896 # HDV max density
VF = 25.0 # Speed free flow
W = 6.25 # Congestion speed
E  = 25.0*0.3 # Speed drop for relaxation

GCAV = 1/(KC*W) # Time headway CAV
GHDV = 1/(KH*W) # Time headway HDV
SCAV = VF/(KC*W)+1/KC #  Desired space headway CAV
SHDV = VF/(KH*W)+1/KH #  Desired space headway HDV

dveh_twy = {'CAV': GCAV, 'HDV': GHDV}
dveh_dwy = {'CAV': 1/KC, 'HDV': 1/KH}

U_MAX = 1.5 # Max. Acceleration
U_MIN = -1.5 # Min. Acceleration
```

Leader static information. This part freezes the leader so reference of the conrol does not change abruptly

```python
stmt = select([traj])
results = connection.execute(stmt).fetchall()
column_names = traj.columns.keys()
trajDf = pd.DataFrame(results, columns = column_names)
network_data = trajDf[(trajDf['ti']==20) & (trajDf['type']=='CAV')]
key = network_data.id.values
ldr = network_data.ldr.values
idx = range(len(key))
dveh_ldr = dict(zip(key,ldr))
dveh_idx = dict(zip(key,idx))
```

Closed loop

```python
%%time
N = 1200 # Simulation steps

t = np.round(np.arange(718)/10,2)
t = [str(t_i) for t_i in t]

t_it = iter(t)
step = iter(range(N))

# Cleaning closed loop tables
stmt = delete(closed)
connection.execute(stmt)
stmt = delete(control)
connection.execute(stmt)

# Query definitions
stmtwriteCL = insert(closed)
stmtwriteUCL = insert(control)

t_ip = 0.0

display(progressSim)

bSuccess = 2
while bSuccess>0:
    try:
        t_i = next(t_it)
        step_i = next(step)

        # Prediction horizon
        t_f = float(t_i) + 5.0
        stmt = select([headway]).where(and_(headway.columns.ti>=t_i,
                                            headway.columns.ti<=t_f))
        h_ref = connection.execute(stmt).fetchall()

        if step_i>=125:
            # Closed looop
            # print(f'Closed loop connected: {t_i, step_i}')
            # print(f'Solved in {n} iterations - Window: {t_i,t_f}')
            stmt = select([closed]).where(and_(closed.columns.ti==t_ip,
                                               closed.columns.type=='CAV'))
            results = connection.execute(stmt).fetchall()

            S, V, DV, U_star, DU, n = compute_control(results, h_ref, 0)

            lVehTrajCL, lVehU = update_state(S, V, DV, U_star, DU, n, results)

        else:
            # Open loop
            # print(f'Open loop connected: {t_i, step_i}')
            stmt = select([traj]).where(and_(traj.columns.ti==t_i,
                                             traj.columns.type=='CAV'))
            results = connection.execute(stmt).fetchall()

            lVehTrajCL, lVehU = format_open_loop(results)

        if lVehTrajCL:
            connection.execute(stmtwriteCL,lVehTrajCL)
            connection.execute(stmtwriteUCL,lVehU)

        t_ip = t_i
        progressSim.value = t_i
    except StopIteration:
        print('Stop by iteration')
        print('Last simulation step at time: {}'.format(t_i))
        bSuccess = 0

```

    FloatProgress(value=5.0, bar_style='info', description='Simulating:', max=71.9)

    Stop by iteration
    Last simulation step at time: 71.7
    CPU times: user 2min 26s, sys: 2.31 s, total: 2min 28s
    Wall time: 2min 36s

```python
refdBDf.to_sql(name='headway', con = engine, if_exists='replace', index=False)
```


# PerfForesightConsumerType: Perfect foresight consumption-saving

[](https://mybinder.org/v2/gh/econ-ark/DemArk/master?filepath=notebooks%2FPerfForesightConsumerType.ipynb)

```python
# Initial imports and notebook setup, click arrow to show
from HARK.ConsumptionSaving.ConsIndShockModel import PerfForesightConsumerType
from HARK.utilities import plotFuncs
from time import clock
import matplotlib.pyplot as plt
import numpy as np
mystr = lambda number : "{:.4f}".format(number)
```

The module $\texttt{HARK.ConsumptionSaving.ConsIndShockModel}$ concerns consumption-saving models with idiosyncratic shocks to (non-capital) income.  All of the models assume CRRA utility with geometric discounting, no bequest motive, and income shocks are fully transitory or fully permanent.

$\texttt{ConsIndShockModel}$ currently includes three models:
1. A very basic "perfect foresight" model with no uncertainty.
2. A model with risk over transitory and permanent income shocks.
3. The model described in (2), with an interest rate for debt that differs from the interest rate for savings.

This notebook provides documentation for the first of these three models.
$\newcommand{\CRRA}{\rho}$
$\newcommand{\DiePrb}{\mathsf{D}}$
$\newcommand{\PermGroFac}{\Gamma}$
$\newcommand{\Rfree}{\mathsf{R}}$
$\newcommand{\DiscFac}{\beta}$

## Statement of perfect foresight consumption-saving model

The $\texttt{PerfForesightConsumerType}$ class the problem of a consumer with Constant Relative Risk Aversion utility
${\CRRA}$
\begin{equation}
U(C) = \frac{C^{1-\CRRA}}{1-\rho},
\end{equation}
has perfect foresight about everything except whether he will die between the end of period $t$ and the beginning of period $t+1$, which occurs with probability $\DiePrb_{t+1}$.  Permanent labor income $P_t$ grows from period $t$ to period $t+1$ by factor $\PermGroFac_{t+1}$.  The consumer faces no artificial borrowing constraint and is able to borrow against the entire future stream of income he will receive.

At the beginning of period $t$, the consumer has an amount of market resources $M_t$ (which includes both market wealth and currrent income) and must choose how much of those resources to consume $C_t$ and how much to retain in a riskless asset $A_t$, which will earn return factor $\Rfree$. The agent's flow of future utility $U(C_{t+n})$ from consumption is geometrically discounted by factor $\DiscFac$ per period. If the consumer dies, he receives zero utility flow for the rest of time.

The agent's problem can be written in Bellman form as:

\begin{eqnarray*}
V_t(M_t,P_t) &=& \max_{C_t}~U(C_t) ~+ \DiscFac (1 - \DiePrb_{t+1}) V_{t+1}(M_{t+1},P_{t+1}), \\
& s.t. & \\
A_t &=& M_t - C_t, \\
M_{t+1} &=& \Rfree A_t + Y_{t+1}, \\
Y_{t+1} &=& P_{t+1}, \\
P_{t+1} &=& \PermGroFac_{t+1} P_t.
\end{eqnarray*}

The consumer's problem is characterized by a coefficient of relative risk aversion $\CRRA$, an intertemporal discount factor $\DiscFac$, an interest factor $\Rfree$, and age-varying sequences of the permanent income growth factor $\PermGroFac_t$ and survival probability $(1 - \DiePrb_t)$.

While it does not reduce the computational complexity of the problem (as permanent income is deterministic, given its initial condition $P_0$), HARK represents this problem with *normalized* variables (represented in lower case), dividing all real variables by permanent income $P_t$ and utility levels by $P_t^{1-\CRRA}$.  The Bellman form of the model thus reduces to:

\begin{eqnarray*}
v_t(m_t) &=& \max_{c_t}~U(c_t) ~+ \DiscFac (1 - \DiePrb_{t+1}) \PermGroFac_{t+1}^{1-\CRRA} v_{t+1}(m_{t+1}), \\
& s.t. & \\
a_t &=& m_t - c_t, \\
m_{t+1} &=& \Rfree/\PermGroFac_{t+1} a_t + 1.
\end{eqnarray*}

## Solution method for PerfForesightConsumerType

Because of the assumptions of CRRA utility, no risk other than mortality, and no artificial borrowing constraint, the problem has a closed form solution.  In fact, the consumption function is perfectly linear, and the value function composed with the inverse utility function is also linear.  The mathematical solution of this model is described in detail in the lecture notes [PerfForesightCRRA](http://econ.jhu.edu/people/ccarroll/public/lecturenotes/consumption/PerfForesightCRRA).

The one period problem for this model is solved by the function $\texttt{solveConsPerfForesight}$, which creates an instance of the class $\texttt{ConsPerfForesightSolver}$. To construct an instance of the class $\texttt{PerfForesightConsumerType}$, several parameters must be passed to its constructor as shown in the table below.

## Example parameter values to construct an instance of PerfForesightConsumerType

| Parameter | Description | Code | Example value | Time-varying? |
| :---: | --- | --- | --- | --- |
| $\DiscFac$ |Intertemporal discount factor  | $\texttt{DiscFac}$ | $0.96$ |  |
| $\CRRA $ |Coefficient of relative risk aversion | $\texttt{CRRA}$ | $2.0$ | |
| $\Rfree$ | Risk free interest factor | $\texttt{Rfree}$ | $1.03$ | |
| $1 - \DiePrb_{t+1}$ |Survival probability | $\texttt{LivPrb}$ | $[0.98]$ | $\surd$ |
|$\PermGroFac_{t+1}$|Permanent income growth factor|$\texttt{PermGroFac}$| $[1.01]$ | $\surd$ |
|$T$| Number of periods in this type's "cycle" |$\texttt{T_cycle}$| $1$ | |
|(none)| Number of times the "cycle" occurs |$\texttt{cycles}$| $0$ | |

Note that the survival probability and income growth factor have time subscripts; likewise, the example values for these parameters are *lists* rather than simply single floats.  This is because those parameters are *time-varying*: their values can depend on which period of the problem the agent is in.  All time-varying parameters *must* be specified as lists, even if the same value occurs in each period for this type.

The last two parameters in the table specify the "nature of time" for this type: the number of (non-terminal) periods in this type's "cycle", and the number of times that the "cycle" occurs.  *Every* subclass of $\texttt{AgentType}$ uses these two code parameters to define the nature of time.  Here, $\texttt{T_cycle}$ has the value $1$, indicating that there is exactly one period in the cycle, while $\texttt{cycles}$ is $0$, indicating that the cycle is repeated in *infinite* number of times-- it is an infinite horizon model, with the same "kind" of period repeated over and over.

In contrast, we could instead specify a life-cycle model by setting $\texttt{T_cycle}$ to $1$, and specifying age-varying sequences of income growth and survival probability.  In all cases, the number of elements in each time-varying parameter should exactly equal $\texttt{T_cycle}$.

The parameter $\texttt{AgentCount}$ specifies how many consumers there are of this *type*-- how many individuals have these exact parameter values and are *ex ante* homogeneous.  This information is not relevant for solving the model, but is needed in order to simulate a population of agents, introducing *ex post* heterogeneity through idiosyncratic shocks.  Of course, simulating a perfect foresight model is quite boring, as there are *no* idiosyncratic shocks other than death!

The cell below defines a dictionary that can be passed to the constructor method for $\texttt{PerfForesightConsumerType}$, with the values from the table here.

```python
PerfForesightDict = {
    # Parameters actually used in the solution method
    "CRRA" : 2.0,                          # Coefficient of relative risk aversion
    "Rfree" : 1.03,                        # Interest factor on assets
    "DiscFac" : 0.96,                      # Default intertemporal discount factor
    "LivPrb" : [0.98],                     # Survival probability
    "PermGroFac" :[1.01],                  # Permanent income growth factor

    # Parameters that characterize the nature of time
    "T_cycle" : 1,                         # Number of periods in the cycle for this agent type
    "cycles" : 0                           # Number of times the cycle occurs (0 --> infinitely repeated)
}
```

## Solving and examining the solution of the perfect foresight model

With the dictionary we have just defined, we can create an instance of $\texttt{PerfForesightConsumerType}$ by passing the dictionary to the class (as if the class were a function).  This instance can then be solved by invoking its $\texttt{solve}$ method.

```python
PFexample = PerfForesightConsumerType(**PerfForesightDict)
PFexample.cycles = 0
PFexample.solve()
```

The $\texttt{solve}$ method fills in the instance's attribute $\texttt{solution}$ as a time-varying list of solutions to each period of the consumer's problem.  In this case, $\texttt{solution}$ will be a list with exactly one instance of the class $\texttt{ConsumerSolution}$, representing the solution to the infinite horizon model we specified.

```python
print(PFexample.solution)
```

    [<HARK.ConsumptionSaving.ConsIndShockModel.ConsumerSolution object at 0x000001EBECEFA3C8>]

Each element of $\texttt{solution}$ has a few attributes. To see all of them, we can use the \texttt{vars} built in function:

the consumption functions reside in the attribute $\texttt{cFunc}$ of each element of $\texttt{ConsumerType.solution}$.  This method creates a (time varying) attribute $\texttt{cFunc}$ that contains a list of consumption functions.

```python
print(vars(PFexample.solution[0]))
```

    {'cFunc': <HARK.interpolation.LinearInterp object at 0x000001EBECEFA198>, 'vFunc': <HARK.ConsumptionSaving.ConsIndShockModel.ValueFunc object at 0x000001EBECEFA550>, 'vPfunc': <HARK.ConsumptionSaving.ConsIndShockModel.MargValueFunc object at 0x000001EBECEFA400>, 'vPPfunc': <HARK.utilities.NullFunc object at 0x000001EBECEFB128>, 'mNrmMin': -50.49994992551661, 'hNrm': 50.49994992551661, 'MPCmin': 0.04428139169919579, 'MPCmax': 0.04428139169919579}

The two most important attributes of a single period solution of this model are the (normalized) consumption function $\texttt{cFunc}$ and the (normalized) value function $\texttt{vFunc}$.  Let's plot those functions near the lower bound of the permissible state space (the attribute $\texttt{mNrmMin}$ tells us the lower bound of $m_t$ where the consumption function is defined).

```python
print('Linear perfect foresight consumption function:')
mMin = PFexample.solution[0].mNrmMin
plotFuncs(PFexample.solution[0].cFunc,mMin,mMin+10.)
```

```python
print('Perfect foresight value function:')
plotFuncs(PFexample.solution[0].vFunc,mMin+0.1,mMin+10.1)
```

An element of $\texttt{solution}$ also includes the (normalized) marginal value function $\texttt{vPfunc}$, and the lower and upper bounds of the marginal propensity to consume (MPC) $\texttt{MPCmin}$ and $\texttt{MPCmax}$.  Note that with a linear consumption function, the MPC is constant, so its lower and upper bound are identical.

## Simulating the perfect foresight consumer model

Suppose we wanted to simulate many consumers who share the parameter values that we passed to $\texttt{PerfForesightConsumerType}$-- an *ex ante* homogeneous *type* of consumers.  To do this, our instance would have to know *how many* agents there are of this type, as well as their initial levels of assets $a_t$ and permanent income $P_t$.

### Setting simulation parameters

Let's fill in this information by passing another dictionary to $\texttt{PFexample}$ with simulation parameters.  The table below lists the parameters that an instance of $\texttt{PerfForesightConsumerType}$ needs in order to successfully simulate its model using the $\texttt{simulate}$ method.

| Description | Code | Example value |
| :---: | --- | --- |
| Number of consumers of this type | $\texttt{AgentCount}$ | $10000$ |
| Number of periods to simulate | $\texttt{T_sim}$ | $120$ |
| Mean of initial log (normalized) assets | $\texttt{aNrmInitMean}$ | $-6.0$ |
| Stdev of initial log  (normalized) assets | $\texttt{aNrmInitStd}$ | $1.0$ |
| Mean of initial log permanent income | $\texttt{pLvlInitMean}$ | $0.0$ |
| Stdev of initial log permanent income | $\texttt{pLvlInitStd}$ | $0.0$ |
| Aggregrate productivity growth factor | $\texttt{PermGroFacAgg}$ | $1.0$ |
| Age after which consumers are automatically killed | $\texttt{T_age}$ | $None$ |

We have specified the model so that initial assets and permanent income are both distributed lognormally, with mean and standard deviation of the underlying normal distributions provided by the user.

The parameter $\texttt{PermGroFacAgg}$ exists for compatibility with more advanced models that employ aggregate productivity shocks; it can simply be set to 1.

In infinite horizon models, it might be useful to prevent agents from living extraordinarily long lives through a fortuitous sequence of mortality shocks.  We have thus provided the option of setting $\texttt{T_age}$ to specify the maximum number of periods that a consumer can live before they are automatically killed (and replaced with a new consumer with initial state drawn from the specified distributions).  This can be turned off by setting it to $\texttt{None}$.

The cell below puts these parameters into a dictionary, then gives them to $\texttt{PFexample}$.  Note that all of these parameters *could* have been passed as part of the original dictionary; we omitted them above for simplicity.

```python
SimulationParams = {
    "AgentCount" : 10000,                  # Number of agents of this type
    "T_sim" : 120,                         # Number of periods to simulate
    "aNrmInitMean" : -6.0,                 # Mean of log initial assets
    "aNrmInitStd"  : 1.0,                  # Standard deviation of log initial assets
    "pLvlInitMean" : 0.0,                  # Mean of log initial permanent income
    "pLvlInitStd"  : 0.0,                  # Standard deviation of log initial permanent income
    "PermGroFacAgg" : 1.0,                 # Aggregate permanent income growth factor
    "T_age" : None,                        # Age after which simulated agents are automatically killed
}

PFexample(**SimulationParams) # This implicitly uses the assignParameters method of AgentType
```

To generate simulated data, we need to specify which variables we want to track the "history" of for this instance.  To do so, we set the $\texttt{track_vars}$ attribute of our $\texttt{PerfForesightConsumerType}$ instance to be a list of strings with the simulation variables we want to track.

In this model, valid elments of $\texttt{track_vars}$ include $\texttt{mNrmNow}$, $\texttt{cNrmNow}$, $\texttt{aNrmNow}$, and $\texttt{pLvlNow}$.  Because this model has no idiosyncratic shocks, our simulated data will be quite boring.

### Generating simulated data

Before simulating, the $\texttt{initializeSim}$ method must be invoked.  This resets our instance back to its initial state, drawing a set of initial $\texttt{aNrmNow}$ and $\texttt{pLvlNow}$ values from the specified distributions and storing them in the attributes $\texttt{aNrmNow_init}$ and $\texttt{pLvlNow_init}$.  It also resets this instance's internal random number generator, so that the same initial states will be set every time $\texttt{initializeSim}$ is called.  In models with non-trivial shocks, this also ensures that the same sequence of shocks will be generated on every simulation run.

Finally, the $\texttt{simulate}$ method can be called.

```python
PFexample.track_vars = ['mNrmNow']
PFexample.initializeSim()
PFexample.simulate()
```

Each simulation variable $\texttt{X}$ named in $\texttt{track_vars}$ will have the *history* of that variable for each agent stored in the attribute $\texttt{X_hist}$ as an array of shape $(\texttt{T_sim},\texttt{AgentCount})$.  To see that the simulation worked as intended, we can plot the mean of $m_t$ in each simulated period:

```python
plt.plot(np.mean(PFexample.mNrmNow_hist,axis=1))
plt.xlabel('Time')
plt.ylabel('Mean normalized market resources')
plt.show()
```

A perfect foresight consumer can borrow against the PDV of his future income-- his human wealth-- and thus as time goes on, our simulated agents approach the (very negative) steady state level of $m_t$ while being steadily replaced with consumers with roughly $m_t=1$.

The slight wiggles in the plotted curve are due to consumers randomly dying and being replaced; their replacement will have an initial state drawn from the distributions specified by the user.  To see the current distribution of ages, we can look at the attribute $\texttt{t_age}$.

```python
N = PFexample.AgentCount
F = np.linspace(0.,1.,N)
plt.plot(np.sort(PFexample.t_age),F)
plt.xlabel('Current age of consumers')
plt.ylabel('Cumulative distribution')
plt.show()
```

The distribution is (discretely) exponential, with a point mass at 120 with consumers who have survived since the beginning of the simulation.

One might wonder why HARK requires users to call $\texttt{initializeSim}$ before calling $\texttt{simulate}$: Why doesn't $\texttt{simulate}$ just call $\texttt{initializeSim}$ as its first step?  We have broken up these two steps so that users can simulate some number of periods, change something in the environment, and then resume the simulation.

When called with no argument, $\texttt{simulate}$ will simulate the model for $\texttt{T_sim}$ periods.  The user can optionally pass an integer specifying the number of periods to simulate (which should not exceed $\texttt{T_sim}$).

In the cell below, we simulate our perfect foresight consumers for 80 periods, then seize a bunch of their assets (dragging their wealth even more negative), then simulate for the reamining 40 periods.

```python
PFexample.initializeSim()
PFexample.simulate(80)
PFexample.aNrmNow += -5. # Adjust all simulated consumers' assets downward by 5
PFexample.simulate(40)

plt.plot(np.mean(PFexample.mNrmNow_hist,axis=1))
plt.xlabel('Time')
plt.ylabel('Mean normalized market resources')
plt.show()
```

```python

```


# Software profesional en Acústica 2020-21 (M2i)

*This notebook contains a modification of the notebook [FEM_Helmholtz_equation](https://github.com/spatialaudio/computational_acoustics/blob/master/FEM_Helmholtz_equation.ipynb), created by Sascha Spors, Frank Schultz, Computational Acoustics Examples, 2018. The text/images are licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/). The code is released under the [MIT license](https://opensource.org/licenses/MIT).*

# Revisiting the numerical solution of the acoustic model: Raviart-Thomas elements and the displacement formulation

This notebook illustrates the numerical solution of the wave equation for harmonic excitation using the so called [Finite Element Method](https://en.wikipedia.org/wiki/Finite_element_method) (FEM). The method aims at an approximate solution by subdividing the area of interest into smaller parts with simpler geometry, linking these parts together and applying methods from the calculus of variations to solve the problem numerically. The FEM is a well established method for the numerical approximation of the solution of partial differential equations (PDEs). The solutions of PDEs are often known analytically only for rather simple geometries. FEM based simulations allow to gain insights into other more complex cases.

## Problem Statement

The acoustic model of a compressible fluid can also be written in terms of the displacement field, given as

\begin{equation}
c^2 \nabla \mathrm{div} \mathbf{u}(\mathbf{x}, t) - \frac{\partial^2}{\partial t^2} \mathbf{u}(\mathbf{x}, t) = \mathbf{0},
\end{equation}

where $\mathbf{u}(\mathbf{x}, t)$ denotes the particle displacement at position $\mathbf{x}$, $c$ the speed of sound.
For an harmonic excitation, we choose the Ansatz $\mathbf{u}(\mathbf{x}, t) = \Re \{ \mathbf{U}(\mathbf{x}, \omega) \mathrm{e}^{\mathrm{j} \omega t} \}$ for the sound pressure with $\omega = 2 \pi f$
Introduction of the complex quantities into the wave equation yields

\begin{equation}
c^2 \nabla \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \mathrm{e}^{\mathrm{j} \omega t} + \omega^2\mathbf{U}(\mathbf{x}, \omega) \mathrm{e}^{\mathrm{j} \omega t} = \mathbf{0},
\end{equation}

and canceling out the $\mathrm{e}^{\mathrm{j} \omega t}$ terms yields the displacement model

\begin{equation}
c^2 \nabla \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) + \omega^2\mathbf{U}(\mathbf{x}, \omega) = \mathbf{0} .
\end{equation}

We aim for a numerical solution of this model in the domain $\Omega$ with respect to a Dirichlet boundary condition
\begin{equation}
\mathbf{U}(\mathbf{x}, \omega)\cdot \mathbf{n} = U_{D} \qquad \text{for } x \in \partial \Omega
\end{equation}

where $P_{D}$ is the Dirichlet data associated with the normal displacement on the boundary.

## Variational Formulation

The FEM is based on expressing the partial differential equation (PDE) to be solved in its variational or weak form.
The first step towards this formulation is to multiply the Helmholtz equation by the test function $\mathbf{V}(\mathbf{x})$

\begin{equation}
c^2 \nabla \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \cdot \mathbf{V}(\mathbf{x})  + \omega^2\mathbf{U}(\mathbf{x}, \omega) \cdot \mathbf{V}(\mathbf{x}) = \mathbf{0} \cdot \mathbf{V}(\mathbf{x}) ,
\end{equation}

followed by integration over the domain $\Omega$

\begin{equation}
\int_\Omega c^2 \nabla \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x  + \omega^2\int_\Omega \mathbf{U} \cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x = \int_\Omega \mathbf{0} \cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x ,
\end{equation}

where $\mathrm{d}x$ denotes a suitably chosen differential element for integration.
Application of a classical Green's formula yields to

\begin{equation}
{-} \int_\Omega \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \mathrm{div} \mathbf{V}(\mathbf{x}) \mathrm{d}x  + \int_{\partial \Omega} \mathbf{V}(\mathbf{x})\cdot \mathbf{n}\ \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \mathrm{d}s + \omega^2 \int_\Omega \mathbf{U}(\mathbf{x}, \omega)\cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x = \int_\Omega \mathbf{0}\cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x.
\end{equation}

This way the differential order of the first integral is lowered which is advisable for application of the FEM.
The second integral vanishes as the variation formulation requires $V(\mathbf{x})\cdot\mathbf{n} = 0$ on $\partial \Omega$ where $(\mathbf{x}, \omega)$ is known - here by the pure Dirichlet boundary condition

This results in the variational/weak formulation of the displacement model

\begin{equation}
{-} \int_\Omega \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \mathrm{div} \mathbf{V}(\mathbf{x}) \mathrm{d}x  + \omega^2 \int_\Omega \mathbf{U}(\mathbf{x}, \omega)\cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x = 0
\end{equation}

It is common to express the integral equation above in terms of the bilinear $a(P, V)$ and linear $L(V)$ forms

\begin{equation}
a(\mathbf{U}, \mathbf{V}) = \int_\Omega \mathrm{div} \mathbf{U}(\mathbf{x}, \omega) \mathrm{div} \mathbf{V}(\mathbf{x}) \mathrm{d}x  - \omega^2 \int_\Omega \mathbf{U}(\mathbf{x}, \omega)\cdot \mathbf{V}(\mathbf{x}) \mathrm{d}x,
\end{equation}

\begin{equation}
L(\mathbf{V}) = 0,
\end{equation}

where

\begin{equation}
a(\mathbf{U}, \mathbf{V}) = L(\mathbf{V}) .
\end{equation}

## Numerical Solution

The numerical solution of the variational problem is based on [FEniCS](https://fenicsproject.org/), an open-source framework for numerical solution of PDEs.
Its high-level Python interface `dolfin` is used in the following to define the problem and compute its solution.
The implementation is based on the variational formulation derived above.
It is common in the FEM to denote the solution of the problem by $u$ and the test function by $v$.
The definition of the problem in FEniCS is very close to the mathematical formulation of the problem.
We limit ourselves to real-valued $\mathbf{U}(\mathbf{x}, \omega)$ due to the assumption of Dirichlet real data.

A function is defined for this purpose, accompanied by a plotting routine for the resulting sound field.

```python
import numpy as np
import dolfin
import mshr
import matplotlib.pyplot as plt
%matplotlib inline

# Compute the Finite Element approximation
def FEM_displacements(mesh, frequency, dirichlet_data, c=343):
    '''numerical solution of the Helmholtz equation using the FEM'''

    # squared wavenumber
    k2 = dolfin.Constant(2*dolfin.pi*frequency/c)**2
    f = dolfin.Expression(("0","0"), degree=1)

    # Define function space (First order Raviart-Thomas elements)
    V = dolfin.FunctionSpace(mesh, "RT", 1) # It's already a vector-valued discrete space!

    # define boundary conditions
    bcs = dolfin.DirichletBC(V, dirichlet_data, "on_boundary")

    # define variational problem
    u = dolfin.TrialFunction(V)
    v = dolfin.TestFunction(V)

    # Define variational problem
    a = k2 * dolfin.inner(u,v)*dolfin.dx - dolfin.div(u)*dolfin.div(v)*dolfin.dx
    L = dolfin.inner(f,v)*dolfin.dx

    # Assemble the linear system of equations
    A, b = dolfin.assemble_system(a, L, bcs)

    # compute solution
    u = dolfin.Function(V)
    dolfin.solve(A, u.vector(), b)

    return u, V

# Compute the Finite Element error
def compute_error(u, uex, V):
    uex_interp = dolfin.interpolate(uex, V)
    return dolfin.errornorm(uex_interp, u)/dolfin.norm(uex_interp)

# Plot the Finite Element approximation
def plot_soundfield(u):
    '''plot solution of FEM-based simulation'''
    fig = plt.figure(figsize=(10,10))
    fig = dolfin.plot(u)
    plt.title(r'$P(\mathbf{x}, \omega)$')
    plt.xlabel(r'$x$ / m')
    plt.ylabel(r'$y$ / m')
    plt.colorbar(fig, fraction=0.038, pad=0.04);
```

### Sound Field in a Rectangular Room with Sound-hard Boundaries

The two-dimensional sound field in a rectangular room whose height is very small compared to the wavelength and with hard boundaries (Dirichlet boundary condition) is computed for the frequency $f=120$ Hz. The Dirichlet data is settled in such a manner that the exact solution is the real part of a plane wave
$$
\mathbf{U}_{ex}(\mathbf{x}, \omega)=(\cos\theta,\sin\theta)\cos(k_{0}x_0+k_{1}x_{1}),\qquad (k_0,k_1)=\frac{\omega}{c}(\cos\theta,\sin\theta),\qquad \theta=\frac{\pi}{4}.
$$

```python
# define geometry and mesh
mesh = dolfin.RectangleMesh(dolfin.Point(0,0), dolfin.Point(5,4), 200, 200, "right/left")

# Physical data
freq = 120.
c = 343.

# Exact solution: plane wave with wavenumber vector (k0,k1)
theta = np.pi/4
k0 = 2*np.pi*freq/c*np.cos(theta)
k1 = 2*np.pi*freq/c*np.sin(theta)
uex = dolfin.Expression(("cos(theta)*cos(k0*x[0]+k1*x[1])","sin(theta)*cos(k0*x[0]+k1*x[1])"), k0=k0, k1=k1, theta=theta, degree=5)

# compute solution
u, V = FEM_displacements(mesh, freq, uex, c)

# compute error
print('L2-relative error:', compute_error(u, uex, V))

# plot sound field
plot_soundfield(u[0])
plt.title(r'$\mathbf{U}_{0}(\mathbf{x}, \omega)$');
plot_soundfield(u[1])
plt.title(r'$\mathbf{U}_{1}(\mathbf{x}, \omega)$');
```

### Study of convergence

Let's check the order of convergence of this finite element discretization computing the error for sucesive refinements of the mesh

```python
# Mesh 50,50
mesh = dolfin.RectangleMesh(dolfin.Point(0,0), dolfin.Point(5,4), 50, 50, "right/left")
u, V = FEM_displacements(mesh, freq, uex, c)
print('Mesh size', mesh.hmax())
print('L2-relative error:', compute_error(u, uex, V))
# Mesh 50,50
mesh = dolfin.RectangleMesh(dolfin.Point(0,0), dolfin.Point(5,4), 100, 100, "right/left")
u, V = FEM_displacements(mesh, freq, uex, c)
print('Mesh size', mesh.hmax())
print('L2-relative error:', compute_error(u, uex, V))
# Mesh 50,50
mesh = dolfin.RectangleMesh(dolfin.Point(0,0), dolfin.Point(5,4), 200, 200, "right/left")
u, V = FEM_displacements(mesh, freq, uex, c)
print('Mesh size', mesh.hmax())
print('L2-relative error:', compute_error(u, uex, V))
```

    Mesh size 0.12806248474865742
    L2-relative error: 0.0025762125662763973
    Mesh size 0.06403124237432907
    L2-relative error: 0.0006451708730117002
    Mesh size 0.032015621187164535
    L2-relative error: 0.0001613628071881609

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT).


# Demo - Helmholtz equation on the unit sphere

**Mikael Mortensen** (email: `mikaem@math.uio.no`), Department of Mathematics, University of Oslo.

Date: **April 20, 2020**

**Summary.** This is a demonstration of how the Python module [shenfun](https://github.com/spectralDNS/shenfun) can be used to solve the
Helmholtz equation on a unit sphere, using spherical
coordinates. This demo is implemented in
a single Python file [sphere_helmholtz.py](https://github.com/spectralDNS/shenfun/blob/master/demo/sphere_helmholtz.py).
If requested the solver will run in parallel using MPI.

Helmholtz on the unit sphere.

## Helmholtz equation

The Helmholtz equation is given as

$$
\begin{equation}
-\nabla^2 u(\boldsymbol{x}) + \alpha u(\boldsymbol{x}) = f(\boldsymbol{x}) \quad \text{for }\, \boldsymbol{x} \in \Omega = \{(x, y, z): x^2+y^2+z^2 = 1\}, \label{eq:helmholtz} \tag{1}
\end{equation}
$$

where $u(\boldsymbol{x})$ is the solution, $f(\boldsymbol{x})$ is a function and $\alpha$ a constant.
We use spherical coordinates $(\theta, \phi)$, defined as

$$
\begin{equation}
 x = r \sin \theta \cos \phi ,
\label{_auto1} \tag{2}
\end{equation}
$$

$$
\begin{equation}
 y = r \sin \theta \sin \phi,
\label{_auto2} \tag{3}
\end{equation}
$$

$$
\begin{equation}
 z = r \cos \theta
\label{_auto3} \tag{4}
\end{equation}
$$

which (with $r=1$) leads to a 2D Cartesian product mesh $(\theta, \phi) \in (0, \pi) \times [0, 2\pi)$
suitable for numerical implementations. There are no boundary
conditions on the problem under consideration.
However, with the chosen Cartesian mesh, periodic
boundary conditions are required for the $\phi$-direction. As such,
the $\phi$-direction will use a Fourier basis $\exp(\imath k \phi)$.

A regular Chebyshev or Legendre basis
$\psi_j(\theta) = \gamma_j(2\theta/\pi-1)$ will be
used for the $\theta$-direction, where $\gamma_j$ could be either
the Chebyshev polynomial of first kind $T_j$ or the Legendre
polynomial $L_j$. Note the mapping from real coordinates $\theta$
to computational coordinates in domain $[-1, 1]$.

The spherical basis functions are as such

$$
v_{jk}(\theta, \phi) = \psi_j(\theta) \exp(\imath k \phi),
$$

and we look for solutions

$$
u(\theta, \phi) = \sum_{j} \sum_{k} \hat{u}_{jk} v_{jk}(\theta, \phi).
$$

A discrete Fourier approximation space with $N$ basis functions is then

$$
V_F^N = \text{span} \{\exp(\imath k \theta)\}, \text{ for } k \in K,
$$

where $K = \{-N/2, -N/2+1, \ldots, N/2-1\}$. For this demo we assume
that the solution is complex, and as such there is no simplification
possible for Hermitian symmetry.

The following approximation space is used for the $\theta$-direction

$$
\begin{equation}
V^N = \text{span} \{\psi_j\}_{j=0}^{N-1},
\label{_auto4} \tag{5}
\end{equation}
$$

and the variational formulation of the problem reads:
find $u \in V^N \otimes V_F^N$ such that

$$
\begin{equation}
   \int_{\Omega} (-\nabla^2 u + \alpha u) v w d\sigma = \int_{\Omega} f v w d\sigma, \quad \forall \, v \in V^N \otimes V_F^N.
\label{eq:u0} \tag{6}
   \end{equation}
$$

Note that integration over the domain is done using
spherical coordinates with an integral measure of $d\sigma=\sin \theta d\theta d\phi$.

## Implementation

A complete implementation is found in the file [sphere_helmholtz.py](https://github.com/spectralDNS/shenfun/blob/master/demo/sphere_helmholtz.py).
Here we give a brief explanation for the implementation. Start by
importing all functionality from [shenfun](https://github.com/spectralDNS/shenfun)
and [sympy](https://sympy.org), where Sympy is required for handeling the
spherical coordinates.

```
from shenfun import *
import sympy as sp

# Define spherical coordinates with unit radius
r = 1
theta, phi = sp.symbols('x,y', real=True, positive=True)
psi = (theta, phi)
rv = (r*sp.sin(theta)*sp.cos(phi), r*sp.sin(theta)*sp.sin(phi), r*sp.cos(theta))
```

Note that the position vector `rv` has three components (for $(x, y, z)$)
even though the computational domain is only 2D.
Also note that Sympy symbols are both positive and real, and $\theta$ is
chosen to be along the first axis and $\phi$ second. This has to agree with
the next step, which is the creation of tensorproductspaces
$V^N \otimes V_F^N$.

```
N, M = 40, 30
L0 = FunctionSpace(N, 'C', domain=(0, np.pi))
F1 = FunctionSpace(M, 'F', dtype='D')
T = TensorProductSpace(comm, (L0, F1), coordinates=(psi, rv, sp.Q.positive(sp.sin(theta))))
```

Spherical coordinates are ensured by feeding `coordinates=(psi, rv, sp.Q.positive(sp.sin(theta)))`
to [TensorProductSpace](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.tensorproductspace.TensorProductSpace), where the restriction `sp.Q.positive(sp.sin(theta))` is there
to help sympy. Operators like [div()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.div),
[grad()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.grad) and  [curl()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.curl) will now work on
items of [Function](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.Function), [TestFunction](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.TestFunction) and
[TrialFunction](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.TrialFunction) using a spherical coordinate system.

To define the equation ([6](#eq:u0)) we first declare
these test- and trialfunctions, and then use code that
is very similar to the mathematics.

```
alpha = 2
v = TestFunction(T)
u = TrialFunction(T)
mats = inner(v, -div(grad(u))+alpha*u)
```

Here `mats` will be a list containing several tensor product
matrices in the form of
[TPMatrix](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.matrixbase.TPMatrix). Since there is only one directions with
non-diagonal matrices ($\theta$-direction) we
can use the generic [SolverGeneric1ND](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.la.SolverGeneric1ND) solver.
Note that some of the non-diagonal matrices will be dense,
which is a weakness of the current method. Also note
that with Legendre one can use integration by parts
instead

```Python
        mats = inner(grad(v), grad(u))
        mats += [inner(v, alpha*u)]
```

To solve the problem we also need to define the function $f(\theta, r)$.
To this end we use sympy and the method of
manufactured solution to define a possible solution `ue`,
and then compute `f` exactly using exact differentiation. We use
the [spherical harmonics function](https://docs.sympy.org/latest/modules/functions/special.html#spherical-harmonics)
to define an analytical solution

```
alpha = 2
sph = sp.functions.special.spherical_harmonics.Ynm
ue = sph(6, 3, theta, phi)

# Compute the right hand side on the quadrature mesh
# That is, compute f = -div(grad(ue)) + alpha*ue
f = (-div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)
fj = Array(T, buffer=f)

# Take scalar product
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)

u_hat = Function(T)
Sol = la.SolverGeneric1ND(mats)
u_hat = Sol(f_hat, u_hat)
```

Having found the solution in spectral space all that is
left is to transform it back to real space.

```
uj = u_hat.backward()
uq = Array(T, buffer=ue)
print('Error =', np.linalg.norm(uj-uq))
```

## Postprocessing
We can refine the solution to make it look better,
and plot on the unit sphere using either [mayavi](https://docs.enthought.com/mayavi/mayavi/)
or [plotly](https://plotly.com/) using the shenfun function
[surf3D()](https://shenfun.readthedocs.io/en/latest/shenfun.utilities.html#shenfun.utilities.surf3D).

```
u_hat2 = u_hat.refine([N*2, M*2])
fig = surf3D(u_hat2.backward().real, wrapaxes=[1])
fig.show()
```


```python
from sympy import chebyshevt, Poly, expand, symbols
from sympy.polys.orthopolys import chebyshevt_poly
import numpy as np
```

```python
dic_list_cheb_2z_1 = []
for i in range(30):
    z = symbols("z")
    dic_poly = Poly(chebyshevt_poly(i, 2 * z - 1), z).as_dict()
    dic_list_cheb_2z_1.append(dic_poly)

dic_list_cheb_z = []
for i in range(30):
    z = symbols("z")
    dic_poly = Poly(chebyshevt_poly(i, z), z).as_dict()
    dic_list_cheb_z.append(dic_poly)

```

```python
np.save('dic_list_cheb_2z_1.npy', dic_list_cheb_2z_1)
np.save('dic_list_cheb_z.npy', dic_list_cheb_z)
```

```python

```


We've been working on a [conference paper](https://github.com/gilbertgede/idetc-2013-paper) to demonstrate the ability to do multibody dynamics with Python. We've been calling this work flow [PyDy](http://pydy.org), short for Python Dynamics. Several pieces of the puzzle have come together lately to really demonstrate the power of the scientific python software packages to handle complex dynamic and controls problems (i.e. IPython notebooks, matplotlib animations, python-control, and our software package mechanics which is a part of SymPy). After writing the draft of our paper, which uses a general n-link pendulum as it's main example, I came across this [blog post by Wolfram](http://blog.wolfram.com/2011/03/01/stabilized-n-link-pendulum/) demonstrating their ability to symbolically derive the equations of motion for the n-link pendulum and stabilize it with an LQR controller. It inspired me to replicate the example as I realized that it was relatively easy to do with all free and open source software!

In this example problem we will derive the equations of motion of an n-link pendulum on a laterally sliding cart and then develop a controller to stabilize it. Balancing a single inverted pendulum is a classic problem that is many times a student's first experience with non-linear dynamics and control. The problem here is extended to a general n-link pendulum and as we will see the equations of motion quickly get messy with greater than 2 links.

The diagram below shows the general description of the problem.

```
from IPython.display import SVG
#SVG(filename='n-pendulum-with-cart.svg')
```

I used these software versions for the following computations:

- IPython: 0.13.1.rc2
- matplotlib: 1.1.1
- NumPy: 1.6.2
- SciPy: 0.10.1
- SymPy: 0.7.2
- python-control: 0.6d

Equations of Motion
===================

We'll start by generating the equations of motion for the system with SymPy **[mechanics](http://docs.sympy.org/dev/modules/physics/mechanics/index.html)**. The functionality that mechanics provides is much more in depth than Mathematica's functionality. In the Mathematica example, Lagrangian mechanics were implemented manually with Mathematica's symbolic functionality. **mechanics** provides an assortment of functions and classes to derive the equations of motion for arbitrarily complex (i.e. configuration constraints, nonholonomic motion constraints, etc) multibody systems in a very natural way. First we import the necessary functionality from SymPy.

```
%matplotlib inline
from sympy import symbols
from sympy.physics.mechanics import *

from sympy import init_printing
#init_printing()
```

Now specify the number of links, $n$. I'll start with 5 since the Wolfram folks only showed four.

```
n = 3
```

**mechanics** will need the generalized coordinates, generalized speeds, and the input force which are all time dependent variables and the bob masses, link lengths, and acceleration due to gravity which are all constants. Time, $t$, is also made available because we will need to differentiate with respect to time.

```
q = dynamicsymbols('q:' + str(n + 1))  # Generalized coordinates
u = dynamicsymbols('u:' + str(n + 1))  # Generalized speeds
f = dynamicsymbols('f')                # Force applied to the cart

m = symbols('m:' + str(n + 1))         # Mass of each bob
l = symbols('l:' + str(n))             # Length of each link
g, t = symbols('g t')                  # Gravity and time
```

Now we can create and inertial reference frame $I$ and define the point, $O$, as the origin.

```
I = ReferenceFrame('I')                # Inertial reference frame
O = Point('O')                         # Origin point
O.set_vel(I, 0)                        # Origin's velocity is zero
```

Secondly, we define the define the first point of the pendulum as a particle which has mass. This point can only move laterally and represents the motion of the "cart".

```
P0 = Point('P0')                       # Hinge point of top link
P0.set_pos(O, q[0] * I.x)              # Set the position of P0
P0.set_vel(I, u[0] * I.x)              # Set the velocity of P0
#P0.set_vel(I, 0)
Pa0 = Particle('Pa0', P0, m[0])        # Define a particle at P0
```

Now we can define the $n$ reference frames, particles, gravitational forces, and kinematical differential equations for each of the pendulum links. This is easily done with a loop.

```
frames = [I]                              # List to hold the n + 1 frames
points = [P0]                             # List to hold the n + 1 points
particles = [Pa0]                         # List to hold the n + 1 particles
forces = [(P0, f * I.x - m[0] * g * I.y)] # List to hold the n + 1 applied forces, including the input force, f
#forces = [(P0, -m[0] * g * I.y)]
kindiffs = [q[0].diff(t) - u[0]]          # List to hold kinematic ODE's

for i in range(n):
    Bi = I.orientnew('B' + str(i), 'Axis', [q[i + 1], I.z])   # Create a new frame
    Bi.set_ang_vel(I, u[i + 1] * I.z)                         # Set angular velocity
    frames.append(Bi)                                         # Add it to the frames list

    Pi = points[-1].locatenew('P' + str(i + 1), l[i] * Bi.x)  # Create a new point
    Pi.v2pt_theory(points[-1], I, Bi)                         # Set the velocity
    points.append(Pi)                                         # Add it to the points list

    Pai = Particle('Pa' + str(i + 1), Pi, m[i + 1])           # Create a new particle
    particles.append(Pai)                                     # Add it to the particles list

    forces.append((Pi, -m[i + 1] * g * I.y))                  # Set the force applied at the point

    kindiffs.append(q[i + 1].diff(t) - u[i + 1])              # Define the kinematic ODE:  dq_i / dt - u_i = 0
```

```
kindiffs
```

    [-u0(t) + Derivative(q0(t), t),
     -u1(t) + Derivative(q1(t), t),
     -u2(t) + Derivative(q2(t), t),
     -u3(t) + Derivative(q3(t), t)]

With all of the necessary point velocities and particle masses defined, the `KanesMethod` class can be used to derive the equations of motion of the system automatically.

```
forces
```

    [(P0, f*I.x - g*m0*I.y), (P1, - g*m1*I.y), (P2, - g*m2*I.y), (P3, - g*m3*I.y)]

```
kane = KanesMethod(I, q_ind=q, u_ind=u, kd_eqs=kindiffs) # Initialize the object
fr, frstar = kane.kanes_equations(forces, particles)     # Generate EoM's fr + frstar = 0
```

The equations of motion are quite long as can been seen below. This is the general nature of most non-simple mutlibody problems. That is why a SymPy is so useful; no more mistakes in algegra, differentiation, or copying in hand written equations.

Simulation
==========

Now that the symbolic equations of motion are available we can simulate the pendulum's motion. We will need some more SymPy functionality and several NumPy functions, and most importantly the integration function from SciPy, `odeint`.

```
from sympy import Dummy, lambdify
from numpy import array, hstack, zeros, linspace, pi
from numpy.linalg import solve
from scipy.integrate import odeint
```

First, define some numeric values for all of the constant parameters in the problem.

```
arm_length = 1. / n                          # The maximum length of the pendulum is 1 meter
bob_mass = 0.01 / n                          # The maximum mass of the bobs is 10 grams
parameters = [g, m[0]]                       # Parameter definitions starting with gravity and the first bob
parameter_vals = [9.81, 0.01 / n]            # Numerical values for the first two
for i in range(n):                           # Then each mass and length
    parameters += [l[i], m[i + 1]]
    parameter_vals += [arm_length, bob_mass]
```

Mathematica has a really nice `NDSolve` function for quickly integrating their symbolic differential equations. We have plans to develop something similar for SymPy but haven't found the development time yet to do it properly. So the next bit isn't as clean as we'd like but you can make use of SymPy's lambdify function to create functions that will evaluate the mass matrix, $M$, and forcing vector, $\bar{f}$ from $M\dot{u} = \bar{f}(q, \dot{q}, u, t)$ as a NumPy function. We make use of dummy symbols to replace the time varying functions in the SymPy equations a simple dummy symbol.

```
dynamic = q + u                                                # Make a list of the states
dynamic.append(f)                                              # Add the input force
dummy_symbols = [Dummy() for i in dynamic]                     # Create a dummy symbol for each variable
dummy_dict = dict(zip(dynamic, dummy_symbols))
kindiff_dict = kane.kindiffdict()                              # Get the solved kinematical differential equations
M = kane.mass_matrix_full.subs(kindiff_dict).subs(dummy_dict)  # Substitute into the mass matrix
F = kane.forcing_full.subs(kindiff_dict).subs(dummy_dict)      # Substitute into the forcing vector
M_func = lambdify(parameters + dummy_symbols, M)               # Create a callable function to evaluate the mass matrix
F_func = lambdify(parameters + dummy_symbols, F)               # Create a callable function to evaluate the forcing vector
```

To integrate the ODE's we need to define a function that returns the derivatives of the states given the current state and time.

```
def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.

    """

    u = 0.0                              # The input force is always zero
    arguments = hstack((args, x, u))     # States, input, and parameters
    dx = np.array(solve(M_func(*arguments), F_func(*arguments))).flatten() # Solving for the derivatives
    return dx

def rhs_partial(x, t, args):
    u = 0
    arg = np.hstack((x, u))
    dx = np.array(solve(mfunc(*arg), ffunc(*arg))).flatten()
    return dx
```

Now that we have the right hand side function, the initial conditions are set such that the pendulum is in the vertical equilibrium and a slight initial rate is set for each speed to ensure the pendulum falls. The equations can then be integrated with SciPy's `odeint` function given a time series.

```
x0 = hstack(( 5, -pi / 2 * ones(len(q) - 1), 0 * 1e-3 * ones(len(u)) )) # Initial conditions, q and u
# note that he then adds the force in the right_hand_side fuction
t = linspace(0, 10, 1000)
#t = linspace(0, .3, 11)# Time vector

xx = x0
#xx = y.T[696]
arguments = np.hstack((parameter_vals, xx, 0))

mm = M_func(*arguments)
ff = F_func(*arguments)

print np.array(solve(mm, ff)).flatten()

import functools
import time

mfunc = functools.partial(M_func, *parameter_vals)
ffunc = functools.partial(F_func, *parameter_vals)

_args = np.hstack((xx, 0))
print np.array(solve(mfunc(*_args), ffunc(*_args))).flatten()
```

    [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
       1.80206776e-15  -7.20827106e-15  -1.56000325e-30   0.00000000e+00]
    [  0.00000000e+00   0.00000000e+00   0.00000000e+00   0.00000000e+00
       1.80206776e-15  -7.20827106e-15  -1.56000325e-30   0.00000000e+00]

```
now = time.time()
y = odeint(right_hand_side, x0, t, args=(parameter_vals,))         # Actual integration
print('Elapsed time: {0:.2f} sec'.format(time.time() - now))

now = time.time()
y = odeint(rhs_partial, x0, t, args=(parameter_vals,))         # Actual integration
print('Elapsed time: {0:.2f} sec'.format(time.time() - now))

# 3 links, 10 sec, 1000 points
# Elapsed time: 1.56 sec
# Elapsed time: 1.28 sec

# 10 links, 10 sec, 1000 points
# Elapsed time: 104.79 sec
# Elapsed time: 88.43 sec
```

    Elapsed time: 0.01 sec
    Elapsed time: 0.01 sec

Plotting
========

The results of the simulation can be plotted with matplotlib.

```
qdeg = np.rad2deg(y.T)
phase = np.rad2deg(np.unwrap(y.T))
```

```
(104 - 88) / 104.
```

    0.15384615384615385

```
q0 = np.deg2rad(qdeg[0])
```

```
fig = plt.figure(figsize=(6, 6))
ax = plt.subplot(111, polar=True)

nn = 100
ax.plot(qdeg[2][:nn], t[:nn], color='r', linewidth=1)
```

```
fig, ax = plt.subplots()

ax.plot(qdeg[1], qdeg[4])
```

```
fig, ax = plt.subplots()

#ax.plot(t, q0)
ax.plot(t, qdeg[-1])
ax.plot(t, phase[-1])
```

```
lab = xlabel('Time [sec]')
leg = legend(dynamic[:y.shape[1] / 2])
```

```
lines = plot(t, y[:, y.shape[1] / 2:])
lab = xlabel('Time [sec]')
leg = legend(dynamic[y.shape[1] / 2:])
```

Animation
=========

matplotlib now includes very nice animation functions for animating matplotlib plots. First we import the necessary functions for creating the animation.

```
from numpy import zeros, cos, sin, arange, around
from matplotlib import pyplot as plt
from matplotlib import animation
from matplotlib.patches import Rectangle
```

The following function was modeled from Jake Vanderplas's [post on matplotlib animations](http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/).

```
def animate_pendulum(t, states, length, filename=None):
    """Animates the n-pendulum and optionally saves it to file.

    Parameters
    ----------
    t : ndarray, shape(m)
        Time array.
    states: ndarray, shape(m,p)
        State time history.
    length: float
        The length of the pendulum links.
    filename: string or None, optional
        If true a movie file will be saved of the animation. This may take some time.

    Returns
    -------
    fig : matplotlib.Figure
        The figure.
    anim : matplotlib.FuncAnimation
        The animation.

    """
    # the number of pendulum bobs
    numpoints = states.shape[1] / 2

    # first set up the figure, the axis, and the plot elements we want to animate
    fig = plt.figure()

    # some dimesions
    cart_width = 0.4
    cart_height = 0.2

    # set the limits based on the motion
    xmin = around(states[:, 0].min() - cart_width / 2.0, 1)
    xmax = around(states[:, 0].max() + cart_width / 2.0, 1)

    # create the axes
    ax = plt.axes(xlim=(xmin, xmax), ylim=(-1.1, 1.1), aspect='equal')

    # display the current time
    time_text = ax.text(0.04, 0.9, '', transform=ax.transAxes)

    # create a rectangular cart
    rect = Rectangle([states[0, 0] - cart_width / 2.0, -cart_height / 2],
        cart_width, cart_height, fill=True, color='red', ec='black')
    ax.add_patch(rect)

    # blank line for the pendulum
    line, = ax.plot([], [], lw=2, marker='o', markersize=6)

    # initialization function: plot the background of each frame
    def init():
        time_text.set_text('')
        rect.set_xy((0.0, 0.0))
        line.set_data([], [])
        return time_text, rect, line,

    # animation function: update the objects
    def animate(i):
        time_text.set_text('time = {:2.2f}'.format(t[i]))
        rect.set_xy((states[i, 0] - cart_width / 2.0, -cart_height / 2))
        x = hstack((states[i, 0], zeros((numpoints - 1))))
        y = zeros((numpoints))
        for j in arange(1, numpoints):
            x[j] = x[j - 1] + length * cos(states[i, j])
            y[j] = y[j - 1] + length * sin(states[i, j])
        line.set_data(x, y)
        return time_text, rect, line,

    # call the animator function
    anim = animation.FuncAnimation(fig, animate, frames=len(t), init_func=init,
            interval=t[-1] / len(t) * 1000, blit=True, repeat=False)

    # save the animation if a filename is given
    if filename is not None:
        anim.save(filename, fps=30, codec='libx264')
```

Now we can create the animation of the pendulum. This animation will show the open loop dynamics.

```
#animate_pendulum(t, y, arm_length)#, filename="open-loop.ogv")
animate_pendulum(t, y, arm_length, filename="open-loop-n3a.mp4")
```

```
from IPython.display import HTML
h = \
"""

"""

h = \
"""

"""

HTML(h)
```

Controller Design
=================

The n-link pendulum can be balanced such that all of the links are inverted above the cart by applying the correct lateral force to the cart. We can design a full state feedback controller based off of a linear model of the pendulum about its upright equilibrium point. We'll start by specifying the equilibrium point and parameters in dictionaries.

```
equilibrium_point = hstack(( 0, pi / 2 * ones(len(q) - 1), zeros(len(u)) ))
equilibrium_dict = dict(zip(q + u, equilibrium_point))
parameter_dict = dict(zip(parameters, parameter_vals))
```

The `KanesMethod` class has method that linearizes the forcing vector about generic state and input perturbation vectors. The equilibrium point and numerical constants can then be substituted in to give the linear system in this form: $M\dot{x}=F_Ax+F_Bu$. The state and input matrices, $A$ and $B$, can then be computed by left side multiplication by the inverse of the mass matrix: $A=M^{-1}F_A$ and $B=M^{-1}F_B$.

```
# symbolically linearize about arbitrary equilibrium
linear_state_matrix, linear_input_matrix, inputs = kane.linearize()
# sub in the equilibrium point and the parameters
f_A_lin = linear_state_matrix.subs(parameter_dict).subs(equilibrium_dict)
f_B_lin = linear_input_matrix.subs(parameter_dict).subs(equilibrium_dict)
m_mat = kane.mass_matrix_full.subs(parameter_dict).subs(equilibrium_dict)
# compute A and B
from numpy import matrix
A = matrix(m_mat.inv() * f_A_lin)
B = matrix(m_mat.inv() * f_B_lin)
```

Now that we have a linear system, the python-control package can be used to design an optimal controller for the system.

```
import control
from numpy import dot, rank
from numpy.linalg import matrix_rank
```

First we can check to see if the system is, in fact, controllable.

```
assert matrix_rank(control.ctrb(A, B)) == A.shape[0]
```

The control matrix is full rank, so now we can compute the optimal gains with a linear quadratic regulator. I chose identity matrices for the weightings for simplicity.

```
K, X, E = control.lqr(A, B, ones(A.shape), 1);
```

The gains can now be used to define the required input during simulation to stabilize the system. The input $u$ is simply the gain vector multiplied by the error in the state vector from the equilibrium point, $u(t)=K(x_{eq} - x(t))$.

```
def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.

    """
    u = dot(K, equilibrium_point - x)    # The controller
    arguments = hstack((x, u, args))     # States, input, and parameters
    dx = array(solve(M_func(*arguments), # Solving for the derivatives
        F_func(*arguments))).T[0]

    return dx
```

Now we can simulate and animate the system to see if the controller works.

```
x0 = hstack(( 0, pi / 2 * ones(len(q) - 1), 1 * ones(len(u)) )) # Initial conditions, q and u
t = linspace(0, 10, 1000)                                       # Time vector
y = odeint(right_hand_side, x0, t, args=(parameter_vals,))      # Actual integration
```

The plots show that we seem to have a stable system.

```
lines = plot(t, y[:, :y.shape[1] / 2])
lab = xlabel('Time [sec]')
leg = legend(dynamic[:y.shape[1] / 2])
```

```
lines = plot(t, y[:, y.shape[1] / 2:])
lab = xlabel('Time [sec]')
leg = legend(dynamic[y.shape[1] / 2:])
```

```
animate_pendulum(t, y, arm_length, filename="closed-loop.ogv")
animate_pendulum(t, y, arm_length, filename="closed-loop.mp4")
```

```
from IPython.display import HTML
h = \
"""

"""
HTML(h)
```

The video clearly shows that our controller can balance all $n$ of the pendulum links. The weightings in the lqr design can be tweaked to give different performance if needed.

This example shows that the free and open source scientific Python tools for dynamics are easily comparable in ability and quality a commercial package such as Mathematica. Besides the current installation hurdles for Python, I'd like to claim that it may better than commercial packages, due to our much more robust SymPy **mechanics** package and the fact that all of the code is liberally licensed for reuse and hacking.

The IPython notebook for this example can be downloaded from https://github.com/gilbertgede/idetc-2013-paper/blob/master/n-pendulum-control.ipynb. Yo ucan try out different $n$ values. I've gotten the equations of motion to compute and an open loop simulation of 10 links. My computer ran out of memory when I tried to compute for $n=50$. The controller weightings and initial conditions will probably have to be adjusted for better performance for $n>5$, but it should work. Let me know the resuls if you play with it.


```python
# Import the SPICE module
import spiceypy
```

```python
# We want to determine the position of our home planet with respect to the Sun.
# The datetime shall be set as "today" (midnight). SPICE requires the Ephemeris Time (ET);
# thus, we need to convert a UTC datetime string to ET.

import datetime

# get today's date
DATE_TODAY = datetime.datetime.today()

# convert the datetime to a string, replacing the time with midnight
DATE_TODAY = DATE_TODAY.strftime('%Y-%m-%dT00:00:00')

# convert the utc midnight string to the corresponding ET
ET_TODAY_MIDNIGHT = spiceypy.utc2et(DATE_TODAY)
```

```python
# oh... an error occurred. The error tells us that a so called "kernel" is missing. These kernels store all
# information that are required for time conversion, pointing, position determination etc.
# For this tutorial the Git repository contains already the necessary kernel. We need to load it first
spiceypy.furnsh('../_kernels/lsk/naif0012.tls')
```

```python
# Let's re-try our first time conversion command
ET_TODAY_MIDNIGHT = spiceypy.utc2et(DATE_TODAY)
```

```python
# It works! How does the value look like?
print(ET_TODAY_MIDNIGHT)
```

    640699269.1855832

```python
# Can we compute now the position and velocity (so called state) of the Earth with respect to the Sun?
# We use the following function to determine the state vector and the so called light time (travel time of the
# light between the Sun and our home planet). Positions are always given in km, velocities in km/s and times in
# seconds

# targ : Object that shall be pointed at
# et : The ET of the computation
# ref : The reference frame. Here, it is ECLIPJ2000 (so Medium article)
# obs :  The observer respectively the center of our state vector computation
EARTH_STATE_WRT_SUN, EARTH_SUN_LIGHT_TIME = spiceypy.spkgeo(targ=399, et=ET_TODAY_MIDNIGHT, \
                                                            ref='ECLIPJ2000', obs=10)
```

```python
# An error occured. Again a kernel error. Well, we need to load a so called spk to load positional information:
spiceypy.furnsh('../_kernels/spk/de432s.bsp')
```

```python
# Let's re-try the computation again
EARTH_STATE_WRT_SUN, EARTH_SUN_LIGHT_TIME = spiceypy.spkgeo(targ=399, et=ET_TODAY_MIDNIGHT, \
                                                            ref='ECLIPJ2000', obs=10)
```

```python
# The state vector is 6 dimensional: x,y,z in km and the corresponding velocities in km/s
print('State vector of the Earth w.r.t. the Sun for "today" (midnight):\n', EARTH_STATE_WRT_SUN)
```

    State vector of the Earth w.r.t. the Sun for "today" (midnight):
     [-1.28760839e+08 -7.76110220e+07  4.32943029e+03  1.48939328e+01
     -2.56361729e+01  1.00712291e-03]

```python
# The (Euclidean) distance should be around 1 AU. Why "around"? Well the Earth revolves the Sun in a slightly
# non-perfect circle (elliptic orbit). First, we compute the distance in km.
import math
EARTH_SUN_DISTANCE = math.sqrt(EARTH_STATE_WRT_SUN[0]**2.0 \
                             + EARTH_STATE_WRT_SUN[1]**2.0 \
                             + EARTH_STATE_WRT_SUN[2]**2.0)
```

```python
# Convert the distance in astronomical units (1 AU)
# Instead of searching for the "most recent" value, we use the default value in SPICE. This way, we can easily
# compare our results with the results of others.
EARTH_SUN_DISTANCE_AU = spiceypy.convrt(EARTH_SUN_DISTANCE, 'km', 'AU')

# Cool, it works!
print('Current distance between the Earth and the Sun in AU:', EARTH_SUN_DISTANCE_AU)
```

    Current distance between the Earth and the Sun in AU: 1.004976583585043

# Orbital speed of the Earth
For this, we need the equation to determine the orbital speed. We assume that the Sun's mass is greater than the mass
of the Earth and we assume that our planet is moving on an almost circular orbit. The orbit velocity $v_{\text{orb}}$ can be approximated with, where $G$ is the gravitational constant, $M$ is the mass of the Sun and $r$ is the distance between the Earth and the Sun:
\begin{align}
    v_{\text{orb}}\approx\sqrt{\frac{GM}{r}}
\end{align}

```python
# First, we compute the actual orbital speed of the Earth around the Sun
EARTH_ORB_SPEED_WRT_SUN = math.sqrt(EARTH_STATE_WRT_SUN[3]**2.0 \
                                  + EARTH_STATE_WRT_SUN[4]**2.0 \
                                  + EARTH_STATE_WRT_SUN[5]**2.0)

# It's around 30 km/s
print('Current orbital speed of the Earth around the Sun in km/s:', EARTH_ORB_SPEED_WRT_SUN)
```

    Current orbital speed of the Earth around the Sun in km/s: 29.648652501451902

```python
# Now let's compute the theoretical expectation. First, we load a pck file that contain miscellanoeus information,
# like the G*M values for different objects

# First, load the kernel
spiceypy.furnsh('../_kernels/pck/gm_de431.tpc')
_, GM_SUN = spiceypy.bodvcd(bodyid=10, item='GM', maxn=1)

# Now compute the orbital speed
V_ORB_FUNC = lambda gm, r: math.sqrt(gm/r)
EARTH_ORB_SPEED_WRT_SUN_THEORY = V_ORB_FUNC(GM_SUN[0], EARTH_SUN_DISTANCE)

# Print the result
print('Theoretical orbital speed of the Earth around the Sun in km/s:', EARTH_ORB_SPEED_WRT_SUN_THEORY)
```

    Theoretical orbital speed of the Earth around the Sun in km/s: 29.710854318668748

```python

```