# Synchornization in complete networks of neurons
Here is the equation of neurons potentails dynamics:
$$
\dot{v_i}=a_i - v_i - \frac{g}{N} \sum_{n|t_n<t} S_{i,l(n)} \delta(t - t_n - t_d)
$$

```python
import os
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
```

## Constants and parameters:

```python
num_neurons = 1000
g = 5
random_input_span = (1.2,2.8)

total_time = 1000
time_step = 0.01
total_steps = int(total_time/time_step)
time_span = np.arange(0,total_time,time_step)

delay_time = 0.1
delay_step = int(delay_time/time_step)

potentail_arr = np.zeros(num_neurons,dtype = float)
random_input = np.random.uniform(*random_input_span,size = num_neurons)
```

```python
spiking_records = np.array( np.zeros((total_steps,num_neurons)), dtype = bool )
```

```python
saving_path_name = 'single_runs'
current_version = 'N{}_T{}_g{}_input_{}_{}'.format(num_neurons,total_time,g,random_input_span[0],random_input_span[1])
try:
    saving_path = os.path.join(saving_path_name, current_version)
    os.makedirs(saving_path)
except:
    pass
```

## Instrumentes to capture spiking frequency
Here we use *cornometer* and number of spikies to measure their frequencies.

```python
spiking_cornometer = np.zeros(num_neurons)
total_spikes_intervals = np.zeros(num_neurons)
total_spikes_num = np.zeros(num_neurons)
```

### Useful functions:

```python
def retarded_spikes_record(step):
    """
    This is the function which take cares of response delay between spikes and their effects.
    """
    if step > delay_step:
        return np.sum(spiking_records[step - delay_step])
    else:
        return 0
```

## Solving the differential equation

```python
for i in tqdm( range(total_steps) ):
    potentail_arr = potentail_arr + (random_input - potentail_arr)*time_step - (g/num_neurons)*retarded_spikes_record(i)
    spiking_records[i] = potentail_arr > 1
    potentail_arr = potentail_arr - 1 * spiking_records[i]
    #frequency calculator
    total_spikes_num = total_spikes_num + spiking_records[i]
    total_spikes_intervals = total_spikes_intervals + spiking_cornometer * spiking_records[i] #add to spike lapsed time
    spiking_cornometer = spiking_cornometer * (spiking_records[i]-1)*(-1) #reset the spikied cornos
    spiking_cornometer = spiking_cornometer + (spiking_records[i]-1)*(-1) #update the not spikied ones
```

    100%|███████████████████████████████████████████████████████████████████████| 100000/100000 [00:03<00:00, 26003.31it/s]

```python

```

```python
plt.hist(potentail_arr)
```

```python

```

# Measuring E (effective field):
we can measure if the total network is synchronized. let's define E as following:

$$
\ddot{E}+ 2\alpha \dot{E}+\alpha^{2}E =\frac{\alpha^{2}}{N} \sum_{n|tـn<t} \delta(t - t_n - t_d)
$$

To do so, we define $M$ as $M = \dot{E} + \alpha E$. As a result we wil have the following systems of equation:
\begin{equation}
    \begin{cases}
      \dot{E} = M - \alpha E\\
      \dot{M} = -  \alpha M + \frac{ \alpha^{2} }{N} \sum_{n|tـn<t} \delta(t - t_n - t_d)
    \end{cases}\,.
\end{equation}

```python
alpha = 20
m_arr = np.zeros(total_steps)
e_arr = np.zeros(total_steps)
```

```python
for i in tqdm(range(0,total_steps-1)):
    m_arr[i+1] = m_arr[i] + time_step*( -alpha*m_arr[i] ) + ( (alpha**2)/num_neurons ) *retarded_spikes_record(i)
    e_arr[i+1] = e_arr[i] + time_step*( m_arr[i] - alpha*e_arr[i] )
```

    100%|█████████████████████████████████████████████████████████████████████████| 99999/99999 [00:01<00:00, 85767.67it/s]

```python
plt.plot(time_span,e_arr)
plt.title('E')
```

```python
plt.plot(e_arr[-1000:])
```

```python
sigma = np.std( e_arr )
sigma
```

    0.23117392007694285

```python

```

# Rasterplot

```python
plt.imshow( spiking_records,aspect= 'auto')
plt.xlabel('Neurons')
plt.ylabel('Time')
plt.title('Rasterplot')
# plt.savefig('raster_plot_N{}.png'.format(num_neurons))
plt.savefig(os.path.join(saving_path, 'raster_plot.png'))
```

```python
spiking_records
```

    array([[False, False, False, ..., False, False, False],
           [False, False, False, ..., False, False, False],
           [False, False, False, ..., False, False, False],
           ...,
           [False, False, False, ..., False, False,  True],
           [False, False, False, ..., False, False, False],
           [False, False, False, ..., False, False, False]])

```python
plt.plot(np.sum( spiking_records, axis = 1))
```

## Calculate frequency

```python
spikies_freq_neurons = total_spikes_intervals * time_step / total_spikes_num
spikies_freq_neurons
```

    C:\ProgramData\Anaconda3\lib\site-packages\ipykernel_launcher.py:1: RuntimeWarning: invalid value encountered in true_divide
      """Entry point for launching an IPython kernel.

    array([ 13.52342466,          nan,          nan,   1.9415625 ,
             1.0230062 ,   1.0433509 ,          nan,          nan,
                    nan,          nan,   1.11962712,          nan,
                    nan,          nan,   2.07605428,   1.4175    ,
                    nan,   3.64285714,          nan,          nan,
                    nan,   1.38467225,   1.32401869,          nan,
                    nan,          nan,          nan,   1.40535411,
                    nan,          nan,          nan,          nan,
             2.89197674,   1.85721495,          nan,          nan,
                    nan,   1.4175    ,          nan,   1.1752669 ,
                    nan,   1.04673362, 229.935     ,   1.04148265,
             1.53826357,   1.41343305,          nan,          nan,
             4.71620853,          nan,          nan,          nan,
             1.29160156,   2.16165217,   4.98625   ,   1.0938232 ,
                    nan,   1.72555556,          nan,          nan,
                    nan,          nan,   1.21114914,          nan,
                    nan,          nan,   1.32759036,   1.39844852,
                    nan,          nan,  27.77484848,   1.68777211,
                    nan,          nan,          nan,   1.04447257,
             1.22320988,          nan,   1.44730321,   1.55770801,
                    nan,   3.51236749,          nan,   2.66630027,
                    nan,          nan,   4.83092233,   2.32697892,
                    nan,   1.2765251 ,          nan,   1.04673362,
             1.39733803,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
             1.39495077,   2.02802041,          nan,          nan,
             1.44655977,   2.02802041,          nan,          nan,
                    nan,   3.57715827,          nan,          nan,
             2.08787815,          nan,   1.5500468 ,          nan,
                    nan,   1.70281304,          nan,   1.06910367,
             1.72555556,          nan,          nan,   2.37591885,
                    nan,   1.15371362,   3.11265625,   2.78120112,
                    nan,   1.39843441,   3.90858824,   1.39495077,
           308.02      ,          nan,          nan,          nan,
                    nan,   1.10574777,   1.41139403,   1.37262794,
             1.91155769,   2.79909859,          nan,          nan,
                    nan,          nan,   2.11016985,          nan,
             1.38467225,   1.36736915,   1.04370917,   2.94352071,
                    nan,          nan,   2.96110119,   1.58370016,
                    nan,          nan,          nan,   2.0827673 ,
                    nan,          nan,   1.82768382,          nan,
             1.0518916 ,          nan,   1.2085    ,          nan,
                    nan,          nan,  10.46893617,   1.31933511,
                    nan,   1.58116242,          nan,   1.21813268,
             4.96139303,   1.1484241 ,          nan,   1.40335219,
                    nan,          nan,          nan,   2.11016985,
                    nan,          nan,   3.29665563,          nan,
             1.04625793,          nan,          nan,          nan,
             2.102537  ,   1.19930993,          nan,   1.4175    ,
             1.17159574,   1.01875514,          nan,   1.42364419,
             7.22934307,          nan,          nan,          nan,
                    nan,   4.30688312,   1.45514663,          nan,
                    nan,          nan,   1.69647863,   1.04625793,
             2.62997354,   1.73882662,          nan,          nan,
                    nan,          nan,          nan,   1.83245387,
                    nan,          nan,          nan,   1.4175    ,
                    nan,   1.39798592,          nan,   2.83805714,
             1.40535411,          nan,   8.86405405,  13.71138889,
                    nan,   2.74765193,          nan,          nan,
                    nan,   1.82904236,   1.51670229, 598.75      ,
             4.45349776,   1.25698352,   2.08043933,   1.68488964,
             1.39297753,          nan,          nan,          nan,
             1.00440609,   2.362019  ,          nan,   1.295     ,
             1.55770801,          nan,          nan,          nan,
                    nan,          nan,   1.26581098,   4.1755042 ,
                    nan,          nan,          nan,          nan,
                    nan,   1.42570402,   1.04558606,          nan,
             1.69988014,          nan,   1.4073617 ,   1.04595565,
                    nan,          nan,          nan,   1.04673362,
             1.24492462,          nan,   1.39599156,          nan,
                    nan,   1.33585464,   2.79120787,          nan,
             1.76800712,          nan,          nan,          nan,
                    nan,   1.39798592,          nan,          nan,
             2.03218814,   4.26029915,   3.51367491,   1.39733803,
             1.77756708,   1.45086257,   1.13189714,          nan,
                    nan,          nan,   2.10701271,   2.11468085,
             2.32697892,          nan,          nan,          nan,
             1.27394602,          nan,   1.40094633,   2.99069069,
                    nan,          nan,  50.92944444,          nan,
                    nan,   1.20225728,   1.53108025,   1.04447257,
                    nan,   2.60918635,          nan,   2.12057569,
                    nan,   2.95228487,          nan,          nan,
             1.1045039 ,          nan,   2.23038117,   1.46886095,
                    nan,   1.73769231,          nan,   2.08043933,
                    nan,          nan,          nan,          nan,
             8.86405405,          nan,   1.39843441,   1.18044048,
                    nan,   2.48184539,          nan,   1.61692182,
             8.33762712,  29.56774194,   5.28446809,          nan,
             1.79696203,          nan,   1.77756708,   1.19973366,
                    nan,          nan,   1.22669554,          nan,
                    nan,          nan,   1.35649795,          nan,
                    nan,          nan,          nan,   2.06184647,
                    nan,          nan,          nan,          nan,
             1.02728216,          nan,          nan,          nan,
                    nan,   1.01004082,   4.17672269,          nan,
                    nan,          nan,   6.69128378,   3.0180303 ,
                    nan,          nan,          nan,          nan,
                    nan,          nan,   1.39733803,   1.85373134,
                    nan,          nan,   1.19535585,   1.0433509 ,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,   1.41140825,
                    nan,          nan,          nan,          nan,
             1.00132591,   1.40335219,   6.87743056,          nan,
             1.20965812,          nan,          nan,   2.5292112 ,
                    nan,          nan,   1.04370917,          nan,
                    nan,   5.67777143,          nan,   1.44730321,
             1.19091346,   1.73464223,          nan,   1.04673362,
             1.50284418,   3.51367491,          nan,          nan,
                    nan,          nan,   1.04559662,          nan,
             1.22725248,   1.32580214,   1.38079277,   1.95793307,
             1.57856916,          nan,          nan,          nan,
                    nan,   1.74303509,          nan,          nan,
                    nan,   1.11706877,   1.20225728,          nan,
                    nan,          nan,          nan,   1.19930993,
             2.88356522,          nan,   2.79120787,          nan,
             5.24921053,   1.1752669 ,   1.39495077, 598.75      ,
                    nan,   1.17525504,          nan,   2.84621777,
                    nan,          nan,   1.39733803,          nan,
            11.18340909,          nan,          nan,          nan,
             2.79120787,          nan,   1.33768194,          nan,
             1.68487267,          nan,          nan,   1.1246311 ,
             1.04819915,   1.40096045,          nan,          nan,
             1.18770983,          nan,   1.1359289 ,          nan,
             3.28574257,   2.01839757,          nan,   1.57856916,
                    nan,          nan,   2.79909859,   1.73769231,
                    nan,   2.15149351,   1.22669554,          nan,
                    nan,   3.21718447,   1.84817844,   1.17525504,
             1.40094633,          nan,          nan,   1.21813268,
             1.77756708,          nan,   1.18341697,   1.39895628,
                    nan,   2.10251586,   2.10868365,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
             1.41139403,          nan,   1.68259322,          nan,
             1.1165389 ,   1.18627545,          nan,   2.688     ,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,   2.77169916,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
             1.25131313,   2.83805714,   2.52274112,          nan,
             2.91750733,          nan,   0.99623364,          nan,
                    nan,          nan,   1.00132591,   1.39733803,
             1.33585464,   3.87809339,   1.26581098,   3.0180303 ,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,   2.55874036,
             1.61692182,          nan,          nan,   1.88339015,
             1.69695726,   2.99978916,   2.28928571,   1.03052029,
                    nan,          nan, 598.75      ,   2.82994302,
           598.75      ,          nan,          nan,   1.70870912,
             2.41385922,          nan,          nan,          nan,
                    nan,          nan,          nan,   2.5921875 ,
                    nan,   1.39844852,          nan,   2.10251586,
                    nan,   1.50284418,   2.88356522,          nan,
             1.30575   ,          nan,   1.39599156,   1.39798592,
                    nan,   1.39798592,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,   2.78120112,          nan,
                    nan,          nan,   2.39633735,   1.18627545,
                    nan,          nan,   1.32329333,   1.39599156,
             2.12057569,          nan,   4.45349776,          nan,
                    nan,          nan,          nan,          nan,
             1.24276942,          nan,  45.8355    ,   2.10251586,
             1.02407446,   2.5921875 ,          nan,   1.38272981,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,   1.19298436,          nan,   1.32938338,
                    nan,          nan,   1.39843441,          nan,
             1.39495077,          nan,          nan,          nan,
                    nan,   4.26029915,   1.39100982,   1.04673362,
                    nan,   2.7400551 ,   1.39599156,          nan,
             1.20225728,          nan,   1.73769231,  11.18340909,
             2.63343915,   2.26828767,          nan,          nan,
                    nan,          nan,   1.04673362,          nan,
                    nan,          nan,          nan,          nan,
             1.05871658,          nan,          nan,   1.65056478,
             2.11468085,   2.11016985,   5.67777143,          nan,
             1.04181053,          nan,   2.29459584,          nan,
                    nan,   1.39495077,          nan,          nan,
                    nan,   1.04707937,          nan,   1.74303509,
                    nan,          nan,   1.24434128,          nan,
                    nan,   1.39599156,   1.14175115,          nan,
                    nan,   2.10251586,          nan,          nan,
             1.90786948,          nan,  10.250625  ,   2.43163814,
                    nan,   7.17688406,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,   1.04673362,          nan,
             1.0085423 ,   1.31756972,          nan,          nan,
             1.04624736,   1.64506623,   1.20859756,   2.14682505,
             1.04447257,          nan,  45.8355    ,          nan,
                    nan,   4.85453659,  83.34545455,   2.91750733,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,          nan,
             1.04625793,          nan,          nan,          nan,
             2.21741071,   4.96139303,          nan,          nan,
            15.79932203,   1.676875  ,          nan,          nan,
                    nan,   3.03954128,          nan,   2.11468085,
             1.06910367,          nan,   2.99972892,          nan,
                    nan,   1.27394602,          nan,          nan,
                    nan,          nan,          nan,          nan,
             6.35471338,          nan,   3.55153571,   1.07027027,
                    nan,          nan,   1.07261105,   1.39599156,
             1.0230062 ,          nan,          nan,   2.10701271,
             2.5921875 , 183.946     ,          nan,   2.393125  ,
             1.69177172,          nan,          nan,          nan,
                    nan,   1.16879717,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,          nan,   1.39495077,
                    nan,   2.02234216,   2.34351415,   1.43674385,
                    nan,   1.73464223,   2.12057569,          nan,
                    nan,   1.04625793,   1.04625793,   1.49600904,
             1.02728216,          nan,   1.45514663,   1.00132591,
                    nan,   6.92559441,          nan,   2.34351415,
             2.79120787,          nan,   1.37262794,          nan,
             1.54638629,   1.01981462,          nan,   1.73160279,
                    nan,          nan,   1.70281304,   3.05830769,
             1.32329333,   1.03270073,          nan,          nan,
                    nan,          nan,          nan,          nan,
                    nan,          nan,   1.40535411,          nan,
             1.39297753,          nan,   1.04739683,   1.04037815,
             2.82994302,   1.57353407,   5.19442708,          nan,
                    nan,          nan,          nan,   2.40799031,
             2.99972892,          nan,          nan,   2.11468085,
                    nan,   3.11265625,          nan,          nan,
             1.04624736,   1.55035937,          nan,   1.04625793,
             1.36313187,          nan,          nan,   4.41150442,
                    nan,          nan,          nan,          nan,
                    nan,   1.46162003,          nan,   1.04559662,
                    nan,          nan,          nan,   1.04447257,
                    nan,          nan,   1.91155769,          nan,
                    nan,          nan,          nan,   3.81850575,
                    nan,          nan,          nan,          nan,
                    nan,          nan,   1.72555556,          nan,
                    nan,          nan,   2.7249589 ,  48.24842105,
             1.75179894,          nan,   1.43674385,   1.27394602,
                    nan,   1.39599156,   1.98325349,   2.38736211,
                    nan,   1.04370917,          nan,          nan,
             1.27160256,   1.39495077,   1.15641774,          nan,
                    nan,   1.04369863,          nan,   2.1759081 ,
                    nan,   1.04673362,          nan,          nan,
             2.08043933, 183.946     ,          nan,   1.0433509 ,
           308.02      ,          nan,          nan,          nan,
                    nan,          nan,   1.3678069 ,          nan,
                    nan,          nan,          nan,          nan,
            38.19458333,   1.74303509,   1.69124361,   2.37591885,
                    nan,          nan,          nan,   1.01348004,
             1.39495077,   1.74303509,          nan,          nan,
                    nan,   2.89197674,          nan,   1.95021569,
             1.04819915,   1.02407446,   1.04673362,   1.04594509])

```python
plt.hist(spikies_freq_neurons)
```

```python

```


Read two saved template Models
Interpolate Ge and Ph data to create input and output targets for training.

1. Load existing template NN models
2. Create parameters for generating training data from source template models
3. Initialize template interpolation (ti) model
4. Train ti model
5. Generate test data for graphing
6. Save ti model to file
7. Generate blender data from ti model output

```python
import torch as torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import os
import numpy as np
import json
import math as m
import random
import time
import NNArt as gy
from sympy.utilities.iterables import multiset_permutations

import plotly as py
import plotly.offline as offline
import plotly.graph_objs as go
import plotly.figure_factory as ff
import plotly.io as pio

CURRENT_MODEL_SAVE_DIR = './save/'

offline.init_notebook_mode(connected=True)
np.set_printoptions(precision=3, suppress=True)
```

```python
class Net(nn.Module):
    def __init__(self, layerSizes, name):
        super(Net, self).__init__()
        self.shape = layerSizes
        self.name = name
        self.seq = nn.Sequential(
            nn.Linear(layerSizes[0],layerSizes[1]),
            nn.ELU(),
            nn.Linear(layerSizes[1],layerSizes[2]),
            nn.ELU(),
            nn.Linear(layerSizes[2],layerSizes[3]),
            nn.ELU()
        )
    def forward(self, x):
        return self.seq(x)
```

```python
#tmModel1
geRange = 1.
batchsize = 10
ge = np.random.random([batchsize,2]) * geRange
tge = torch.Tensor(ge)
tm0Out = tmModel1(tge).data.numpy()

tmModel1.shape, tge.shape, tm0Out.shape
```

    ([2, 10, 15, 21], torch.Size([10, 2]), (10, 21))

```python
def genRcTrain(relModel,genModel, parms):
    epochs = parms['epochs']
    lrR = parms['lrR']
    lrG = parms['lrG']
    tModels = parms['tModels']
    device = parms['device']
    numSamples = parms['numSamples']
    rFilter = parms['rFilter']
    numContextDims=1 #extra emergent dimensions

    minLossG = 1000
    maxLossG = -10
    minLossR = 10
    maxLossR = 0
    criterion = torch.nn.MSELoss(reduction='sum')
    relOptimizer = torch.optim.Adam(relModel.parameters(), lr=lrR)
    genOptimizer = torch.optim.Adam(genModel.parameters(), lr=lrG)
    for epoch in range(epochs):
        #Each epoch get a new training set based on random input to the template models
        relTarget, genTarget, _ = gy.createGenRcTargetsFromModelList(tModels, numSamples, numContextDims)

        genOptimizer.zero_grad()
        relOptimizer.zero_grad()

        relOutput = relModel(genTarget)
        genOutput = genModel(relOutput)

        lossG = criterion(genOutput, genTarget)
        lossR = criterion(relOutput*rFilter, relTarget*rFilter)

        (lossG).backward(retain_graph=True)
        (lossR).backward()

        genOptimizer.step()
        relOptimizer.step()

        minLossG = min(lossG.data.item(), minLossG)
        maxLossG = max(lossG.data.item(), maxLossG)
        minLossR = min(lossR.data.item(), minLossR)
        maxLossR = max(lossR.data.item(), maxLossR)

        if (epoch) % (epochs/10) == 0:
            print("Epoch {0:,.0f} - lossR: {1:,.5f} - lossG: {2:,.5f}".format(
                epoch, lossR/numSamples, lossG/numSamples))

    print("Epoch {0:,.0f} - lossR: {1:,.5f}:{2:,.5f} - lossG: {3:,.5f}:{4:,.5f}".format(epoch+1,
                minLossR/numSamples, maxLossR/numSamples, minLossG/numSamples, maxLossG/numSamples))

    return
```

```python
'''
1. Load existing template NN models
templateModel2

'''

tmModel0 = Net([2,10,15,21],'tmModel0')
tmModel1 = Net([2,10,15,21],'tmModel1')

PATH=CURRENT_MODEL_SAVE_DIR+'templateModel0.pth'
tmModel0.load_state_dict(torch.load(PATH))
PATH=CURRENT_MODEL_SAVE_DIR+'templateModel1.pth'
tmModel1.load_state_dict(torch.load(PATH))
tmModel0.eval()
tmModel1.eval()

tmModel0.shape, tmModel1.shape,PATH
```

    ([2, 10, 15, 21], [2, 10, 15, 21], './save/templateModel1.pth')

```python
'''
2. Initialize P-Gc-P

'''

#initialize network
manualSeed = 1

np.random.seed(manualSeed)
torch.manual_seed(manualSeed)

relModel = Net([21,15,10,3],'relModel')
genModel = Net([3,10,15,21],'genModel')
relModel.train()
genModel.train()

# device = torch.device("cuda:0")
device = torch.device("cpu")
if device:
    relModel.to(device)
    genModel.to(device)
    tmModel0.to(device)
    tmModel1.to(device)
    print("Transfer to use ", device)

relModel.shape, genModel.shape
```

    Transfer to use  cpu

    ([21, 15, 10, 3], [3, 10, 15, 21])

```python
'''
3. Train template interpolation model

'''
start = time.time()

parms={ 'epochs':100000,'lrR':.0001,'lrG':.0002,'numSamples':1,'rFilter':torch.Tensor([1,1,.05]),
       'tModels':[tmModel0,tmModel1],'device':device}
genRcTrain(relModel,genModel,parms)

print('time: ', time.time()-start)
relModel.train = False
genModel.train = False
```

    Epoch 0 - lossR: 1.84756 - lossG: 3.13059
    Epoch 10,000 - lossR: 0.19716 - lossG: 0.01196
    Epoch 20,000 - lossR: 0.00103 - lossG: 0.00510
    Epoch 30,000 - lossR: 0.06058 - lossG: 0.00220
    Epoch 40,000 - lossR: 0.00432 - lossG: 0.00054
    Epoch 50,000 - lossR: 0.00259 - lossG: 0.00171
    Epoch 60,000 - lossR: 0.00032 - lossG: 0.00095
    Epoch 70,000 - lossR: 0.00100 - lossG: 0.00310
    Epoch 80,000 - lossR: 0.00284 - lossG: 0.00219
    Epoch 90,000 - lossR: 0.00256 - lossG: 0.00107
    Epoch 100,000 - lossR: 0.00000:3.46201 - lossG: 0.00002:6.14693
    time:  242.9731731414795

```python
'''
parms={ 'epochs':100000,'lrG':.0004,'lrP':.0005,'numSamples':1,'gFilter':torch.Tensor([1,1,.05]),
Epoch 100,000 - lossG: 0.00000:0.27628 - lossP: 0.00002:0.07341
parms={ 'epochs':100000,'lrG':.0002,'lrP':.0003,'numSamples':1,'gFilter':torch.Tensor([1,1,.05]),
Epoch 100,000 - lossG: 0.00000:0.23505 - lossP: 0.00001:0.06539
parms={ 'epochs':100000,'lrG':.0001,'lrP':.0002,'numSamples':1,'gFilter':torch.Tensor([1,1,.05]),
Epoch 100,000 - lossG: 0.00000:0.21284 - lossP: 0.00001:0.04594
Epoch 100,000 - lossG: 0.00000:0.21077 - lossP: 0.00001:0.03137
Epoch 100,000 - lossG: 0.00000:0.19727 - lossP: 0.00001:0.03196
'''
1==1
```

```python
'''
4. Review Emergent Context

'''
morphIndices=[0.,0.1,.2,.3,.5,.7,.9,1.0, 1.1]
timeIndices=[0.,0.1,.2,.3,.5,.7,.9,1.0, 1.1]

#sample source templates to get ph samples
relIn = gy.createRmInput(morphIndices, timeIndices, sPermutation=[])
phTout0 = tmModel0(torch.Tensor(relIn))
phTout1 = tmModel1(torch.Tensor(relIn))

#use template generated ph samples to feed through contex models
relCout0 = relModel(phTout0)
relCout1 = relModel(phTout1)

genOut0 = genModel(relCout0)
genOut1 = genModel(relCout1)

# gin.shape, pout2.shape, geCout2.shape, phCout2.shape
```

```python
g20 = np.array([relIn[0,:,1],relCout0[0,:,2].detach().numpy()])
g22 = np.array([relIn[2,:,1],relCout0[2,:,2].detach().numpy()])
g24 = np.array([relIn[4,:,1],relCout0[4,:,2].detach().numpy()])
g26 = np.array([relIn[6,:,1],relCout0[6,:,2].detach().numpy()])
g27 = np.array([relIn[7,:,1],relCout0[7,:,2].detach().numpy()])
g28 = np.array([relIn[8,:,1],relCout0[8,:,2].detach().numpy()])

g30 = np.array([relIn[0,:,1],relCout1[0,:,2].detach().numpy()])
g32 = np.array([relIn[2,:,1],relCout1[2,:,2].detach().numpy()])
g34 = np.array([relIn[4,:,1],relCout1[4,:,2].detach().numpy()])
g36 = np.array([relIn[6,:,1],relCout1[6,:,2].detach().numpy()])
g37 = np.array([relIn[7,:,1],relCout1[7,:,2].detach().numpy()])
g38 = np.array([relIn[8,:,1],relCout1[8,:,2].detach().numpy()])

gx = [g20,g30,g22,g32,g24,g34,g26,g36,g28,g38]

```

```python
layout = go.Layout(
    title = 'Initial Curves',
    height = 700,
    width = 800,
    hovermode= 'closest',
    xaxis= dict(
        title= 'X',
        range = [-.1,1.5],
    ),
    yaxis=dict(
        title= 'Y',
        range = [-.6,1.],
    )
)
layout.title = 'PhC changes into GeOutput Sp Space for GeInput: '
f = gy.graphCurves(gx, layout, labels=['s2t0', 's3t0','s2t2', 's3t2','s2t4', 's3t4','s2t6', 's3t6','s2t8', 's3t8', ], dashRange=gy.rangeList([range(2,4),range(6,8)]))

offline.iplot(f)

```

```python

i = 5
ph0 = phTout0.data.numpy()[i]
ph1 = phTout1.data.numpy()[i]

p01_0 = genOut0.data.numpy()[i]
p01_1 = genOut1.data.numpy()[i]

# nx = np.concatenate((p0,p01_0))
nx = np.concatenate((p1,p01_1))

gx = gy.FormatNetworkToGraphColumns(nx, xOffset=0, yOffset=14, colPoints=7)

layout = go.Layout(
    title = 'Initial Curves',
    height = 600,
    width = 600,
    hovermode= 'closest',
    xaxis= dict(
        title= 'X',
        range = [-.4,1.2],
    ),
    yaxis=dict(
        title= 'Y',
        range = [-.4,1.2],
    )
)
f = gy.graphCurves(gx, layout, dashRange=range(0,9))
offline.iplot(f)

```

```python

i = 7
ph0 = phTout0.data.numpy()[i]
ph1 = phTout1.data.numpy()[i]

p01_0 = genOut0.data.numpy()[i]
p01_1 = genOut1.data.numpy()[i]

# nx = np.concatenate((p0,p01_0))
nx = np.concatenate((p1,p01_1))

gx = gy.FormatNetworkToGraphColumns(nx, xOffset=0, yOffset=7, colPoints=7)

layout = go.Layout(
    title = 'Initial Curves',
    height = 600,
    width = 600,
    hovermode= 'closest',
    xaxis= dict(
        title= 'X',
        range = [-.3,1.1],
    ),
    yaxis=dict(
        title= 'Y',
        range = [-.2,.2],
    )
)
f = gy.graphCurves(gx, layout, dashRange=range(0,9))
offline.iplot(f)

```

```python

i = 7
ph0 = phTout0.data.numpy()[i]
ph1 = phTout1.data.numpy()[i]

p01_0 = genOut0.data.numpy()[i]
p01_1 = genOut1.data.numpy()[i]

# nx = np.concatenate((p0,p01_0))
nx = np.concatenate((p1,p01_1))

# gx = gy.FormatNetworkToGraphColumns(nx, xOffset=0, yOffset=7, colPoints=7)
# gx = gy.FormatNetworkToGraphColumns(nx, xOffset=0, yOffset=14, colPoints=7)
gx = gy.FormatNetworkToGraphColumns(nx, xOffset=7, yOffset=14, colPoints=7)

layout = go.Layout(
    title = 'Initial Curves',
    height = 600,
    width = 600,
    hovermode= 'closest',
    xaxis= dict(
        title= 'X',
        range = [-.5,.5],
    ),
    yaxis=dict(
        title= 'Y',
        range = [-.8,.2],
    )
)
f = gy.graphCurves(gx, layout, dashRange=range(0,9))
offline.iplot(f)

```

```python
'''
5. Graph data
'''
ti = 7
nx = np.array(phTout1[ti].data.numpy())
gx = gy.FormatNetworkToGraphColumns(nx, xOffset=0, yOffset=7, colPoints=7)

layout = go.Layout(
    title = 'Initial Curves',
    height = 600,
    width = 600,
    hovermode= 'closest',
    xaxis= dict(
        title= 'X',
        range = [-.4,1.1],
    ),
    yaxis=dict(
        title= 'Y',
        range = [-.25,.25],
    )
)

layout.title = 'Time 0 for different species combinations'
f = gy.graphCurves(gx, layout)
offline.iplot(f)
```

```python
'''
7. Save two template interpolation model to file

'''

# PATH = CURRENT_MODEL_SAVE_DIR +'GEN PhRcPh-T0T1-mask 11p05.pth'
# torch.save(phModel.state_dict(), PATH)
# PATH = CURRENT_MODEL_SAVE_DIR +'REL PhRc-T0T1-mask 11p05.pth'
# torch.save(geModel.state_dict(), PATH)

```

    '\n7. Save two template interpolation model to file\n\n'

```python
'''
9. Generate blender data from template interpolation model output

#phOut.shape torch.Size([10, 7, 23])

'''

# phOut = phModel(torch.Tensor(torch.Tensor(inpt0)))

# x = phOut.detach().numpy()
# fb = gy.FormatNetworkToBlender(x, xOffset=0, yOffset=7, zOffset=14, colPoints=7)

# with open(CURRENT_MODEL_SAVE_DIR +'btest.json', 'w') as fp:
#     json.dump(fb.tolist(),fp)

# x.shape, fb.shape
```


 # Gomes, & Michaelides (2003)

 # "[Portfolio choice with internal habit formation: a life-cycle model with uninsurable labor income risk](https://www.sciencedirect.com/science/article/pii/S1094202503000590)"

 - Notebook modified by Tung-Sheng Hsieh based on [Mateo Velásquez-Giraldo](https://github.com/Mv77) and [Matthew Zahn](https://sites.google.com/view/matthew-v-zahn/matthew-v-zahn).

This notebook uses the [Econ-ARK/HARK](https://github.com/econ-ark/hark) toolkit to describe the main results and reproduce the figures in the linked paper. The main HARK tool used here is the $\texttt{PortfolioConsumerType}$ class. For an introduction to this module, see the [HARK's Documentation](https://hark.readthedocs.io/en/latest/).

I thank Chris Carroll, Tao Wang, Hsin-Yi Hung and Mateo Velásquez-Giraldo for comments and guidance.

```python
# This cell does some preliminary set up

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

# Import relevenat HARK tools
import HARK.ConsumptionSaving.ConsPortfolioModel as cpm

# This is a jupytext paired notebook that autogenerates BufferStockTheory.py
# which can be executed from a terminal command line via "ipython BufferStockTheory.py"
# But a terminal does not permit inline figures, so we need to test jupyter vs terminal
# Google "how can I check if code is executed in the ipython notebook"
from IPython import get_ipython # In case it was run from python instead of ipython

# If the ipython process contains 'terminal' assume not in a notebook
def in_ipynb():
    try:
        if 'terminal' in str(type(get_ipython())):
            return False
        else:
            return True
    except NameError:
        return False

# Determine whether to make the figures inline (for spyder or jupyter)
# vs whatever is the automatic setting that will apply if run from the terminal
if in_ipynb():
    # %matplotlib inline generates a syntax error when run from the shell
    # so do this instead
    get_ipython().run_line_magic('matplotlib', 'inline')
else:
    get_ipython().run_line_magic('matplotlib', 'auto')
```

### The model

In the paper the author extend the finite horizon portfolio choice model, in which consumer has uninsurable labor income risk, cannot borrow or short-sell, and has a fixed costs to hold risky assets (stocks). The extension introduces internal habit formation in preferences. The habit process is specified as a function of the household's past consumption, and both the ratio and additive specification are studied.

#### The optimization problem of ratio habit model
\begin{equation}
\max _{\left\{S_{i t}, B_{i t}\right\}_{t=1}^{T}} E_{1} \sum_{t=1}^{T} \beta^{t-1}\left\{\prod_{j=0}^{t-1} p_{j}\right\} U\left(C_{i t}, H_{i t}\right)
\end{equation}

where the felicity function $U\left(C_{it},H_{it}\right)=\frac{\left\{C_{t} / H_{t}^{\gamma}\right\}^{1-\rho}}{1-\rho}$. The parameter $\gamma \in [0,1]$  controls the importance of habit, and the felicity function becomes the standard time separable function with CRRA parameter equal to $\rho$ when $\gamma = 0$

subject to

\begin{array}{l}H_{i t}=(1-\lambda) H_{i t-1}+\lambda C_{i t-1} \\X_{i t+1}=S_{i t} R_{t+1}+B_{i t} R_{f}+Y_{i t+1}-F \cdot I_{P} \cdot P_{i t+1} \\X_{i t}=S_{i t}+B_{i t}+C_{i t} \\B_{i t} \geqslant 0, \quad S_{i t} \geqslant 0 \\R_{t+1}-R_{f}=\mu+\varepsilon_{t+1}, \quad \varepsilon_{t} \sim N\left(0, \sigma_{\varepsilon}^{2}\right) \\Y_{i t}=P_{i t} U_{i t}, \quad P_{i t}=\exp \left(f\left(t, Z_{i t}\right)\right) P_{i t-1} N_{i t}, \quad \text { if } t \leqslant K \\Y_{i t}=\theta P_{i K}, \quad \text { if } t>K\end{array}

#### Habit evolving process
\begin{equation}
H_{i t}=(1-\lambda) H_{i t-1}+\lambda C_{i t-1}
\end{equation}

$\lambda \in [0,1]$, denotes the weight of habit formation of the consumption in the last period.
#### Labor income process

\begin{equation}
Y_{i t}=P_{i t} U_{i t}, \quad \text { if } t \leqslant K\\
Y_{i t}=\theta P_{i K}, \quad \text { if } t>K
\end{equation}

where $f(t,Z_{it})$ is a deterministic function of age and household characteristics $Z_{i,t}$. $P_{it}$ is the permanent component of income, and $U_{it}$ is a transitory part of it. The log of transitory shock, $\ln U_{it}$ is assuemed to be independend and identically distributed to $(-0.5\sigma^2_u,\sigma^2_u)$ while the log of permanent income shock $\ln N_{it} \sim (-0.5\sigma_n^2,\sigma_n^2)$

$P_{it}$ follows the following process:
\begin{equation}
P_{i t}=e^{\left(f\left(t, Z_{i t}\right)\right)} P_{i t-1} N_{i t}
\end{equation}

meaning that the log of $P_{it}$ evolves as a random walk with a deterministic drift, $f(t,Z_{it})$

$K=46$ is assumed to be the exogenous and deterministic retiment age at 65, and $\theta\in [0,1]$ is the replacement ratio.

#### Asset and wealth accumulation

There are two assets, one risk-free and one risky, available for consumers to allocate their savings.

- T-bills: after-tax risk-free return $R^f$.

- Stocks: $R_{t+1}-R^f = \mu + \varepsilon_{t+1}$, where $\mu$ is the mean equity premium and $\varepsilon_t \sim N(0,\sigma^2_\varepsilon)$

Denote the next period's cash on hand as $X_{i,t+1}$. When the stock market entry fee has not been paid yet, we have

\begin{equation}
X_{i,t+1}=S_{it}R_{t+1}+B_{it}R^f+Y_{i,t+1}-FI_pP_{i,t+1}
\end{equation}

$I_p$ is a dummy variable which takes value 1 when the cost is incurred for the first-time entrance, and takes 0 otherwise. $F$ is the fixed cost, and the total entree fee is scaled by timing the permanent income $P_{i,t+1}$

##### Intra-temporal allocation decision
\begin{equation}
X_{it}=C_{it}+S_{it}+B_{it}
\end{equation}

Where $C_{it}$ denotes the consumption expenditure, $S_{it}$ denotes the amount of stock holdings, and $B_{it}$ denotes the amount of reskless asset holdings.

##### Borrowing and short-sell constraint
\begin{equation}
B_{it}\geq 0\\
S_{it}\geq 0
\end{equation}

#### Matching labor income in HARK

In HARK's consumption-saving models, the income process takes the form
\begin{equation}
    \ln Y_t = \ln P_t + \ln \theta_t
\end{equation}
where $P_t$ represents permanent income and $\ln \theta_t \sim N(0,\sigma_\theta)$ transitory shocks to income. Permanent income evolves according to
\begin{equation}
    \ln P_{t+1} = \ln \Gamma_{t+1} +  \ln \psi_{t+1} + \ln P_t\\
    \ln P_{t} = \ln \Gamma_{t} +  \ln \psi_{t} + \ln P_{t-1}
\end{equation}
where $\Gamma_{t+1}$ is a deterministic growth factor, and $\ln \psi_{t+1} \sim N(0,\sigma_\psi)$ a permanent income shock.

To represent the author's assumptions in HARK, we express both income processes as sums of deterministic components and i.i.d shocks
\begin{align}
\text{GM}           &\quad& \ln Y_{i,t} &=& f(t,Z_{i,t}) + \ln P_{i,0} &+&\sum_{k=1}^t \ln N_{i,k} &+&\qquad \ln U_{i,t}\\
\text{Cocco et. al} &\quad& \ln Y_{i,t} &=& f(t,Z_{i,t}) + v_{i,0} &+&\sum_{k=1}^t u_{i,k} &+&\qquad \varepsilon_{i,t} \\
\text{HARK}         &\quad& \ln Y_{i,t} &=& \sum_{k=1}^t \ln \Gamma_k + \ln P_{i,0} &+&\sum_{k=1}^t \ln \psi_{i,k} &+& \qquad \ln \theta_{i,t}
\end{align}

These representations make evident the mapping that we use

|caps meanings| HARK |shocks_dist.| Cocco et. al |shocks_dist.| GM (2003)       |shocks_dist                   |
| :----: | :---: | :-----------: |:-----------:|:-----------:|:-----------:|:-----------------:|
|period zero permanent income| $\ln$ $P_{i,0}$ || $f(0,Z_{i,0})$ + $v_{i,0}$ | |$f(0,Z_{i,0})$ + $\ln$ $P_{i,0}$ ||
|period k+1 deterministic income grwoth factor| $\ln$ $\Gamma_{k+1}$||$f(k+1$, $Z_{i,k+1})$ - $f(k,Z_{i,k})$||$f(k+1$, $Z_{i,k+1})$ - $f(k,Z_{i,k})$||
|period k permanent income shock|$\ln$ $\psi_{i,k}$| $N(0,\sigma_\psi^2)$ |$u_{i,k}$ | $N(0,\sigma_u^2)$ |$\ln$ $N_{i,k}$|$N(-0.5\sigma^2_n,\sigma_n^2)$|
|period k transitory income shock|$\ln$ $\theta_{i,t}$| $N(0,\sigma_\theta^2)$ |$\varepsilon_{i,t}$|$N(0,\sigma_\varepsilon^2)$ |$\ln$ $U_{i,t}$|$N(-0.5\sigma^2_u,\sigma_u^2)$|

and to achieve a retirement income that is equal to a fraction $\lambda$ of permanent income in the last working period $K$, we simply make $\Gamma_{K+1} = \theta$ and $\Gamma_{t} = 1$ $\forall t>K+1$.

In the Appendix B the author specified that shocks are noramally distributed, and that follows the HARK assumption that we have the expected level of transitory component equals 1.

### Calibration—Summary

__1. Preferences and life cycle__

 | Parameter | Description | Code | Value |
 |:---:| ---         | ---  | :---: |
 | $\beta$ | Time Preference Factor | $\texttt{DiscFac}$ | 0.95 |
 | $\rho$ | Coefficient of Relative Risk Aversion| $\texttt{CRRA}$ | 2 |
 | $p_t$ | Survival Probability | $\texttt{LivPrb}$ | [0.6809,0.99845] |
 | $t_0$ | Starting Age | $\texttt{t_start}$ | 19 |
 | $t_r$ | Retirement Age | $\texttt{t_ret}$ | 65 |
 | $T$ | Maximum Age | $\texttt{t_end}$ | 100 |
 | $\gamma$ | habit importance | $\texttt{Hgamma}$ | 0.8 |
 | $\lambda$ | habit "catch-up" speed | $\texttt{Hlambda}$ | 0.5 |

Two new parameters habit importance and habit "catch-up" speed are added and are set to 0.8 and 0.5 respectively as benchmark.
The paper did not specify how they determine the living probability, but following the rule above, I use the
living probability from fotran code of Cocco et al. (2005), and add the first period living probability artificially since in Cocco et al. (2005) the agent starts to work since age 20 while in this paper the agent starts to work since age 19.

__2. Income process and the finanical assets__

 | Parameter | Description | Code | Value|
 |:---:| ---         | ---  | :---: |
 | $f(t,Z_{i,t})$| Average income at each stage of life | $\texttt{det_income}$ | $ \exp($0.530339 + 0.16818 $t$+ (0.0323371/10) $t^2$ + (0.0019704/100) $t^3)$ |
 | $\lambda$ | Last Period Labor Income Share for Retirement | $\texttt{repl_fac}$ | 0.68212 |
 | $\log \Gamma$ | Permanent Income Growth Factor | $\texttt{PermGroFac}$ | $\{\log f_{t+1} - \log f_t\}^{t_r+1}_{t=20}$ |
 | $\mathsf{R}$ | Interest Factor | $\texttt{Rfree}$ | 1.02 |
 | $\mu$ | Average Stock Return | $\texttt{RiskyDstnFunc}$ \& $\texttt{RiskyDrawFunc}$ | 1.06 |
 | $\sigma_\eta$ | Std Dev of Stock Returns | $\texttt{RiskyDstnFunc}$ \& $\texttt{RiskyDrawFunc}$ | 0.18 |

The deterministic labor income profile is chosen to reflect the hump shape of earnings over the life cycle and the parameter values are taken from Cocco et al. (1999). The retirement transfers are also calibrated using the estimation results from Cocco et al. (1999) and are set to around 68 percent of labor income in the last period of working life. These calibration values were also used in Cocco et al. (2005) and are used here.
__3. Shocks__

 | Parameter | Description | Code | Value |
 |:---:| ---         | ---  | :---: |
 | $\sigma_n$ | Std Dev of Log Permanent Shock| $\texttt{PermShkStd}$ | 0.08 |
 | $\sigma_u$ | Std Dev of Log Transitory Shock| $\texttt{TranShkStd}$ | 0.1 |

Authors used 0.08 and 0.1 respectively for the standard deviation of log permanent shock and log transitory shock, which are consistent with the estimation by Carroll (1997).

For reference, the Cocco et al.'s source Fortran code that includes these paramerization details is available on [Gomes' personal page](http://faculty.london.edu/fgomes/research.html).

```python
# Set up the parameters in the paper
# Relative risk aversion
CRRA = 2
# Discount factor
DiscFac = 0.95

time_params = {'Age_born': 19, 'Age_retire': 65, 'Age_death': 100}
t_start = time_params['Age_born']
t_ret   = time_params['Age_retire'] # We are currently interpreting this as the last period of work
t_end   = time_params['Age_death']

# Survival probabilities from the author's Fortran code
n = 80
survprob = np.zeros(n+1)
survprob[0] = 0.99854
survprob[1] = 0.99845
survprob[2] = 0.99839
survprob[3] = 0.99833
survprob[4] = 0.9983
survprob[5] = 0.99827
survprob[6] = 0.99826
survprob[7] = 0.99824
survprob[8] = 0.9982
survprob[9] = 0.99813
survprob[10] = 0.99804
survprob[11] = 0.99795
survprob[12] = 0.99785
survprob[13] = 0.99776
survprob[14] = 0.99766
survprob[15] = 0.99755
survprob[16] = 0.99743
survprob[17] = 0.9973
survprob[18] = 0.99718
survprob[19] = 0.99707
survprob[20] = 0.99696
survprob[21] = 0.99685
survprob[22] = 0.99672
survprob[23] = 0.99656
survprob[24] = 0.99635
survprob[25] = 0.9961
survprob[26] = 0.99579
survprob[27] = 0.99543
survprob[28] = 0.99504
survprob[29] = 0.99463
survprob[30] = 0.9942
survprob[31] = 0.9937
survprob[32] = 0.99311
survprob[33] = 0.99245
survprob[34] = 0.99172
survprob[35] = 0.99091
survprob[36] = 0.99005
survprob[37] = 0.98911
survprob[38] = 0.98803
survprob[39] = 0.9868
survprob[40] = 0.98545
survprob[41] = 0.98409
survprob[42] = 0.9827
survprob[43] = 0.98123
survprob[44] = 0.97961
survprob[45] = 0.97786
survprob[46] = 0.97603
survprob[47] = 0.97414
survprob[48] = 0.97207
survprob[49] = 0.9697
survprob[50] = 0.96699
survprob[51] = 0.96393
survprob[52] = 0.96055
survprob[53] = 0.9569
survprob[54] = 0.9531
survprob[55] = 0.94921
survprob[56] = 0.94508
survprob[57] = 0.94057
survprob[58] = 0.9357
survprob[59] = 0.93031
survprob[60] = 0.92424
survprob[61] = 0.91717
survprob[62] = 0.90922
survprob[63] = 0.90089
survprob[64] = 0.89282
survprob[65] = 0.88503
survprob[66] = 0.87622
survprob[67] = 0.86576
survprob[68] = 0.8544
survprob[69] = 0.8423
survprob[70] = 0.82942
survprob[71] = 0.8154
survprob[72] = 0.80002
survprob[73] = 0.78404
survprob[74] = 0.76842
survprob[75] = 0.75382
survprob[76] = 0.73996
survprob[77] = 0.72464
survprob[78] = 0.71057
survprob[79] = 0.6961
survprob[80] = 0.6809

# Fix indexing problem (fortran starts at 1, python at 0)
# survprob = np.delete(survprob, [0])

# Labor income

# They assume its a polinomial of age. Here are the coefficients
a=-2.170042+2.700381
b1=0.16818
b2=-0.0323371/10
b3=0.0019704/100

time_params = {'Age_born': 19, 'Age_retire': 65, 'Age_death': 100}
t_start = time_params['Age_born']
t_ret   = time_params['Age_retire'] # We are currently interpreting this as the last period of work
t_end   = time_params['Age_death']

# They assume retirement income is a fraction of labor income in the
# last working period
repl_fac = 0.68212

# Compute average income at each point in (working) life
f = np.arange(t_start, t_ret+1,1)
f = a + b1*f + b2*(f**2) + b3*(f**3)
det_work_inc = np.exp(f)

# Retirement income
det_ret_inc = repl_fac*det_work_inc[-1]*np.ones(t_end - t_ret)

# Get a full vector of the deterministic part of income
det_income = np.concatenate((det_work_inc, det_ret_inc))

# ln Gamma_t+1 = f_t+1 - f_t = ln Y_it+1 - ln Y_it
gr_fac = np.exp(np.diff(np.log(det_income)))

# Now we have growth factors for T_end-1 periods.

# Finally define the normalization factor used by CGM, for plots.
norm_factor_0 = det_income

# %% Shocks

# Transitory and permanent shock variance from the paper
std_tran_shock = 0.1
std_perm_shock = 0.08

# Vectorize. (HARK turns off these shocks after T_retirement)
std_tran_vec = np.array([std_tran_shock]*(t_end-t_start))
std_perm_vec = np.array([std_perm_shock]*(t_end-t_start))

# %% Financial instruments

# Risk-free factor
Rfree = 1.02

# Creation of risky asset return distributions

Mu = 0.04 # Equity premium
Std = 0.18 # standard deviation of rate-of-return shocks

RiskyAvg = Mu + Rfree
RiskyStd = Std
# Make a dictionary to specify the rest of params
dict_portfolio={}
dict_portfolio['CRRA'] = CRRA
dict_portfolio['Rfree'] = Rfree
dict_portfolio['DiscFac'] = DiscFac
dict_portfolio['T_age'] = t_end-t_start+1
dict_portfolio['T_cycle'] = t_end-t_start
dict_portfolio['T_retire'] = t_ret-t_start+1
dict_portfolio['cycles'] = 1
dict_portfolio['PermShkStd'] = std_perm_vec
# dict_portfolio['PermShkCount'] = 3
dict_portfolio['PermShkCount'] = 7 # Gaussian quadrature with J=7
dict_portfolio['TranShkStd'] = std_tran_vec
dict_portfolio['RiskyAvg'] = RiskyAvg
dict_portfolio['RiskyStd'] = RiskyStd
dict_portfolio['RiskyShareCount'] = 30
dict_portfolio['LivPrb'] = survprob.tolist()
dict_portfolio['PermGroFac'] = gr_fac.tolist()
# dict_portfolio['TranShkCount'] = 3
dict_portfolio['TranShkCount'] = 7 # Gaussian quadrature with J=7
dict_portfolio['UnempPrb']= 0
dict_portfolio['UnempPrbRet']= 0
dict_portfolio['IncUnemp']= 0
dict_portfolio['IncUnempRet']=  0
dict_portfolio['BoroCnstArt']=  0
dict_portfolio['tax_rate']= 0.0
# dict_portfolio['RiskyCount']=3
dict_portfolio['RiskyCount']=7
dict_portfolio['RiskyShareCount']= 30
dict_portfolio['vFuncBool']=False
dict_portfolio['CubicBool']= False
# Simulation params
dict_portfolio['AgentCount']= 10
dict_portfolio['pLvlInitMean']=  np.log(det_income[0]) # Mean of log initial permanent income (only matters for simulation)
dict_portfolio['pLvlInitStd']=  std_perm_shock  # Standard deviation of log initial permanent income (only matters for simulation)
dict_portfolio['T_sim']= (t_end - t_start+1)*50
# Unused params required for simulation
dict_portfolio['PermGroFacAgg']=  1
dict_portfolio['aNrmInitMean']=  -50.0 # Agents start with 0 assets (this is log-mean)
dict_portfolio['aNrmInitStd' ]=  0.0
age_plot_params = [20, 30, 55, 75]
```

## Key Results
### Consumption with different habit states
The author first shows that for a given habit level, the shape of the consumption policy rule is the same as in the standard
buffer stock saving literature. Consumption equals cash on hand below a certain threshold level, and beyond that level the marginal propensity to consume rapidly falls. The same threshold effect is also replicated below.
To make this point clearer, we can actually write out the marginal propensity to consume with habit formation under perfect foresight to get some insight. According to [PerfForesightConsumerType.ipynb](https://github.com/econ-ark/HARK/blob/master/examples/ConsIndShockModel/PerfForesightConsumerType.ipynb) and [Carroll (2019)](https://www.econ2.jhu.edu/people/ccarroll/papers/BufferStockTheory/#MPCnvrs), we can write the recursive form of MPC as
\begin{align}
\kappa_{t}^{-1} & = & 1 + \kappa_{t+1}^{-1}(\Rfree \beta)^{1/\rho}/G
\end{align}
The term that times the inverse of MPC in t+1 period on the right hand side is defined in this [lecture note](http://www.econ2.jhu.edu/people/ccarroll/public/LectureNotes/Consumption/PerfForesightCRRA.pdf) as "Growth Patience Factor". By some simple derivation, we can find that the GPF with habit can be written as the following:
\begin{align}
\frac{(\Rfree \beta)^{\frac{1}{\rho}}}{G}\left(\frac{h_{t+1}}{h_t}\right)^{\gamma(1-\frac{1}{\rho})}
\end{align}
showing that the habit stock is affecting the MPC of consumer, and this result is also used to construct a new HARK consumer type with under perfect foresight with habits.

Secondly, the author showed that the threshold is increasing in the stock of habits since a higher habit level requires additional consumption to maintain the same level of utility. This is also replicated below.
The volatile characteristics of the consumption function was not reproduced, and may require more fundamental rewrite of the coding inlucding defning new HARK consumer type.

```python
### Figure 1: See how different habit states change the consumption

# The habit parameters
Hgamma = 0.8
Hlambda = 0.5

# Run the model
dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist() # habit catch-up speed was incorporated in the growth factor
agent = cpm.PortfolioConsumerType(**dict_portfolio)
agent.solve()

# Create a grid of market resources for the plots

mMin = 0    # Minimum ratio of assets to income to plot
mMax = 10 # Maximum ratio of assets to income to plot
mPts = 100 # Number of points to plot

eevalgrid = np.linspace(0,mMax,mPts) # range of values of assets for the plot

### Habit
# Plot different habit states at age 25- Figure 1
plt.figure()
# Rather than generating different states of habits, we simply suppose different initial habit levels to check the qualitative results.
for H0 in [0.2,0.4,0.6]:

    H = np.ones(81)*H0 # generates habit stocks for every life period
    H = H0**(Hgamma*(1-CRRA)) # define the factor that will be used in normalization

    norm_factor = norm_factor_0/H # define the factor that will be used in normalization

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].cFuncAdj(eevalgrid/norm_factor[a-age_born])*norm_factor[a-age_born],
             label = 'H0 = {:.2f}'.format(H0))
plt.xlabel('Wealth')
plt.ylabel('Consumption')
plt.title('Consumption Policy Function with deifferent habit level at age 25')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

We can see that lower habits generate a higher level of saving than higher habit levels.

## Consumption with different habit importance and catch-up speed
Authors showed in figure 2 that the consumption function with mean habit level associated with the simulations of the model decreases faster than the one without the habits. The key finding here is that the consumption with habit decreases faster, and is reproduced using the habit stock different than that used by the author. The change in habit stock actually has the same effect when altering the importance factor of habit $\gamma$, and the results is showed in figure 4 that consumption function with higher importance decreases faster, consistent with the findings in figure 2.

```python
### Figure 2 and 4: See how Hgamma and Hlambda change the consumption

plt.figure()
for Hgamma in [0, 0.5, 0.8]: # set the different level of habit importance from low to high
    H0=0.2 # use 0.2 as the initial habit level
    H = np.ones(81)*H0
    H = H0**(Hgamma*(1-CRRA))

    norm_factor = norm_factor_0/H

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].cFuncAdj(eevalgrid/norm_factor[a-age_born])*norm_factor[a-age_born],
             label = 'Hgamma = {:.2f}'.format(Hgamma))
plt.xlabel('Wealth')
plt.ylabel('Consumption')
plt.title('Consumption Policy Function')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

## Share of wealth in stock
The author argued that the presence of future labor income increases the demand for stocks, but this needs to be compared to the share function without income shocks and cannot be inferred simply from figure 3.
In addition, a higher habit importance has in fact the same effect as the high habit stocks, showing that lower habit importance generates a higher level of saving than higher habit imprtance. The result is consistent with figure 1.
The following replication shows the same order of share function with different habit importance, but differ with the figure in the initial range of cash-on-hand. The author's figure shows initial range of cash-on-hand for which positive savings take place for the lowest habit states, thus generating positive stock holdings, while no savings occur under the highest states.
The replication successfully replicated that, for high levels of liquid wealth a higher habit level is associated with a higher share. Based on the inference before, a higher habit level is associated with lower saving rate, generating a smaller  ratio of savings to future labor income for a given level of cash on hand. Therefore they will invest larger proportion of their savings to risky assets.

```python
### Figure 3: See how different habit states change the share of wealth in stock.

# The habit parameters
Hgamma = 0.8
Hlambda = 0.5

# Run the model
dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
agent = cpm.PortfolioConsumerType(**dict_portfolio)
agent.solve()

# Create a grid of market resources for the plots

mMin = 0    # Minimum ratio of assets to income to plot
mMax = 50 # Maximum ratio of assets to income to plot
mPts = 100 # Number of points to plot

eevalgrid = np.linspace(0,mMax,mPts) # range of values of assets for the plot

# Plot share of wealth with different habit states- Figure 3
plt.figure()
for H0 in [0.2,0.4,0.6, 0.8]:

    H = np.ones(81)*H0
    H = H0**(Hgamma*(1-CRRA))

    norm_factor = norm_factor_0/H

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].ShareFuncAdj(eevalgrid/norm_factor[a-age_born]),
             label = 'H0 = {:.2f}'.format(H0))
plt.xlabel('Wealth')
plt.ylabel('Risky portfolio share')
plt.title('Risky portfolio policy function')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

## Insensitivity
Here we test the claimed the insensitivity of catch-up speed to consumption, and insensitivity of habit importance and catch-up speed to share function.
The replicatin with change in Hgamma is insentive to share function, but the replication with change in Hlambda does not replicate the results, due to the techniques I used. I incorporate the Hlambda with the income growth factor, and thus changing the Hlambda in fact changed the present value of income and will affect the agent's allocation decision.

```python
### Try to reproduce the result that the catch-up speed Hlambda is not very important for consumption function- page 739's description

Hgamma = 0.8
H0=0.2
H = np.ones(82)*H0
H = H0**(Hgamma*(1-CRRA))

norm_factor = norm_factor_0/H

mMin = 0    # Minimum ratio of assets to income to plot
mMax = 50 # Maximum ratio of assets to income to plot
mPts = 100 # Number of points to plot

eevalgrid = np.linspace(0,mMax,mPts) # range of values of assets for the plot

plt.figure()
for Hlambda in [0.5, 0.6, 0.7, 0.8]:

    # Run the model
    dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
    agent = cpm.PortfolioConsumerType(**dict_portfolio)
    agent.solve()

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].cFuncAdj(eevalgrid/norm_factor[a-age_born])*norm_factor[a-age_born],
             label = 'Hlambda ={:.2f}'.format(Hlambda))
plt.xlabel('Wealth')
plt.ylabel('Consumption')
plt.title('Consumption Policy Function under different $\lambda$')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

```python
### Try to reproduce the result that the catch-up speed Hlambda is not very important for share function- page 739's description

Hgamma = 0.8
H0=0.2
H = np.ones(82)*H0
H = H0**(Hgamma*(1-CRRA))

norm_factor = norm_factor_0/H

mMin = 0    # Minimum ratio of assets to income to plot
mMax = 50 # Maximum ratio of assets to income to plot
mPts = 100 # Number of points to plot

eevalgrid = np.linspace(0,mMax,mPts) # range of values of assets for the plot

plt.figure()
for Hlambda in [0.5, 0.6, 0.7, 0.8]:

    # Run the model
    dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
    agent = cpm.PortfolioConsumerType(**dict_portfolio)
    agent.solve()

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].ShareFuncAdj(eevalgrid/norm_factor[a-age_born]),
             label = 'Hlambda = {:.2f}'.format(Hlambda))
plt.xlabel('Wealth')
plt.ylabel('Risky portfolio share')
plt.title('Risky portfolio policy function')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

```python
### Try to reproduce the result that the habit importance Hgamma does not really affect share function- page 739's description

Hlambda = 0.5
H0=0.2
H = np.ones(82)*H0
H = H0**(Hgamma*(1-CRRA))

norm_factor = norm_factor_0/H

mMin = 0    # Minimum ratio of assets to income to plot
mMax = 50 # Maximum ratio of assets to income to plot
mPts = 100 # Number of points to plot

eevalgrid = np.linspace(0,mMax,mPts) # range of values of assets for the plot

plt.figure()
for Hgamma in [0.5, 0.6, 0.7, 0.8]:

    # Run the model
    dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
    agent = cpm.PortfolioConsumerType(**dict_portfolio)
    agent.solve()

    # Ages
    ages = [25]
    age_born = time_params['Age_born']

    a = 25
    plt.plot(eevalgrid,
             agent.solution[a-age_born].ShareFuncAdj(eevalgrid/norm_factor[a-age_born]),
             label = 'Hgamma = {:.2f}'.format(Hgamma))
plt.xlabel('Wealth')
plt.ylabel('Risky portfolio share')
plt.title('Risky portfolio policy function')
plt.legend()
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

Here the replication shows that the share is not sensitive to habit importance since it does not affect risk attitude.

## Constant Habit
Given the paper is right that Hlambda does not affect both consumption and risky share decision, to simplify the programming, we can actually let Hlambda = 0, meaning that the habit is constant at a certain level throught the agent's whole life cycle, and get the same qualitative or even quantitative result.

### Simulations

Using the policy functions obtained from solving the model we present a series of simulations to highlight features of the model.

We first run a few simulations to verify the quality of our calibration.

The figures below show simulated levels of permanent income and risky portfolio shares for 50 agents over their life spans. We can see the model generates a heterogeneous permanent income distribution. Interestingly, we have the same couterfactual prediction of the
model: households invest almost all of their wealth in stocks throughout the life-cycle (and therefore stock market participation takes place as soon as positive saving is undertaken)

```python
# Set up simulation parameters

Hlambda = 0 # Constant Habit
# Rerun the model
dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
agent = cpm.PortfolioConsumerType(**dict_portfolio)
agent.solve()

# Number of agents and periods in the simulation.
agent.AgentCount = 50 # Number of instances of the class to be simulated.
# Since agents can die, they are replaced by a new agent whenever they do.

# Number of periods to be simulated
agent.T_sim = 80

# Set up the variables we want to keep track of.
agent.track_vars = ['aNrmNow','cNrmNow', 'pLvlNow', 't_age', 'ShareNow','mNrmNow']

# Run the simulations
agent.initializeSim()
agent.simulate()

# Present diagnostic plots.
plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.pLvlNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Permanent income')
plt.title('Simulated Income Paths')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.mNrmNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Normalized wealth')
plt.title('Simulated Wealth Accumulation')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.ShareNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Risky share')
plt.title('Simulated Risky Portfolio Shares')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

#### The average life cycle patterns

We now compare the behavior of the mean values of variables of interest at different ages, conditional on survival.
The plot below illustrates the average dynamics across all of the simulated agents of market resources and risky shares with habit that replicate figures 7 and figure 8. The plot follows the general pattern observed in the original paper.
The average outcomes without habits requires a more fundamental change in the solution techniques used here.

```python
# Number of agents and periods in the simulation.
agent.AgentCount = 50 # Number of instances of the class to be simulated.
# Since agents can die, they are replaced by a new agent whenever they do.

# Number of periods to be simulated
agent.T_sim = 80*50

# Run the simulations
agent.initializeSim()
agent.simulate()

raw_data = {'Age': agent.t_age_hist.flatten()+time_params['Age_born'],
            'pIncome': agent.pLvlNow_hist.flatten(),
            'rShare': agent.ShareNow_hist.flatten(),
            'nrmM': agent.mNrmNow_hist.flatten(),
            'nrmC': agent.cNrmNow_hist.flatten()}

Data = pd.DataFrame(raw_data)
Data['Cons'] = Data.nrmC * Data.pIncome
Data['M'] = Data.nrmM * Data.pIncome

# Find the mean of each variable at every age
AgeMeans = Data.groupby(['Age']).mean().reset_index()
plt.figure()
plt.plot(AgeMeans.Age, AgeMeans.M,
         label = 'Market resources')
plt.legend()
plt.xlabel('Age')
plt.title('Variable Means Conditional on Survival')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot

plt.figure()
plt.plot(AgeMeans.Age, AgeMeans.rShare, label = 'Mean')
plt.legend()

axes = plt.gca()
axes.set_xlim([20,100])
axes.set_ylim([0,1.5])

plt.xlabel('Age')
plt.ylabel('Risky Share')
plt.title('Risky Portfolio Share Mean Conditional on Survival')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

### Other results in the original paper

#### Results with fixed participation cost

#### Additive habit model

### Conclusion

The authors conclude that the inclusion of habit in the life cycle choice model does not help to explain the empirical evidence of low stock holdings in assets. With internal habit, the intertemporal consumption smoothing becomes more important and households increase their wealth accumulation. As a result, they have a stronger incentive to pay the stock market entry cost and start investing in equities very early in life. Moreover, after paying the fixed cost, they invest virtually all of their wealth in stocks.
### Comparison with Cocco et al. (2005)
This article does not provide explanation for how the optimal rsiky asset decreases with wealth, which is addressed by Cocco et al. (2005). In fact, the paper imply a counterfactual almost 100% stock market participation throught the whole life shown in the simulation.
Both paper shown that the optimal risky asset share decreases in wealth. Both paper also agree that labor income increases risky asset holdings, as it is seen as a closer substitute of risk-free assets.
Another reason for the discrepancy may come form the different calibration parameters use in two papers. The following is the robustness check with several interesting parameters.

part worthwhile to mention that, in calibration parameters, two papers use quite different magnitude of CRRA coefficients, which are 10 in this paper and 2 in Cocco et al. (2005). The low risk aversion used in this paper's calibration may be the reason why the simulated portfolio share is quiet constant at 1 throughout the whole life.

#### Risk aversion
set back the risk aversion from low (=2) to high (=10 in Cocco et al.). The replication below reproduce the similar share function as shown in CGMPortfolio.ipynb. This make sense since if the agent is not sensitive to risk, it will be optimal to put more share of wealth in the stock markets.

```python
# Set up simulation parameters
#std_tran_shock = np.sqrt(0.0738)
#std_perm_shock = np.sqrt(0.0106)

Hlambda = 0 # Constant Habit
CRRA = 10 # Cocco et al.
dict_portfolio['CRRA'] = CRRA
# Rerun the model
dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
agent = cpm.PortfolioConsumerType(**dict_portfolio)
agent.solve()

# Number of agents and periods in the simulation.
agent.AgentCount = 50 # Number of instances of the class to be simulated.
# Since agents can die, they are replaced by a new agent whenever they do.

# Number of periods to be simulated
agent.T_sim = 80

# Set up the variables we want to keep track of.
agent.track_vars = ['aNrmNow','cNrmNow', 'pLvlNow', 't_age', 'ShareNow','mNrmNow']

# Run the simulations
agent.initializeSim()
agent.simulate()

# Present diagnostic plots.
plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.pLvlNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Permanent income')
plt.title('Simulated Income Paths')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.mNrmNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Normalized wealth')
plt.title('Simulated Wealth Accumulation')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.ShareNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Risky share')
plt.title('Simulated Risky Portfolio Shares')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

#### Income distribution
Two papers also use different distribution of permanent income. Since both paper agree that permanent income is more like the substitute to risk free asset, it is expected that the distribution of income may also affect their share in wealth. The following robustness check simulation revert permanent and transitory shocks back to Cocco's level, maintaining CRRA =2, but the share functin does not change much. This may suggest that the change in CRRA has dominant effect.

```python
# Set up simulation parameters
std_tran_shock = np.sqrt(0.0738)
std_perm_shock = np.sqrt(0.0106)

# Vectorize. (HARK turns off these shocks after T_retirement)
std_tran_vec = np.array([std_tran_shock]*(t_end-t_start))
std_perm_vec = np.array([std_perm_shock]*(t_end-t_start))

dict_portfolio['PermShkStd'] = std_perm_vec
dict_portfolio['TranShkStd'] = std_tran_vec

Hlambda = 0 # Constant Habit
dict_portfolio['CRRA'] = 2 # set CRRA back to 2
# Rerun the model
dict_portfolio['PermGroFac'] = (gr_fac*(1-Hlambda)).tolist()
agent = cpm.PortfolioConsumerType(**dict_portfolio)
agent.solve()

# Number of agents and periods in the simulation.
agent.AgentCount = 50 # Number of instances of the class to be simulated.
# Since agents can die, they are replaced by a new agent whenever they do.

# Number of periods to be simulated
agent.T_sim = 80

# Set up the variables we want to keep track of.
agent.track_vars = ['aNrmNow','cNrmNow', 'pLvlNow', 't_age', 'ShareNow','mNrmNow']

# Run the simulations
agent.initializeSim()
agent.simulate()

# Present diagnostic plots.
plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.pLvlNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Permanent income')
plt.title('Simulated Income Paths')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.mNrmNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Normalized wealth')
plt.title('Simulated Wealth Accumulation')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)

plt.figure()
plt.plot(agent.t_age_hist+time_params['Age_born'], agent.ShareNow_hist,'.')
plt.xlabel('Age')
plt.ylabel('Risky share')
plt.title('Simulated Risky Portfolio Shares')
plt.grid()

if not in_ipynb():
    plt.show(block=False)
    plt.pause(1)
else:
    plt.show(block=True)
```

```python
np.sqrt(0.0738)
```

    0.2716615541441225

### Bibliographic entry of the original article

```python
### Bibtex entry
'''
@Article{RePEc:red:issued:v:6:y:2003:i:4:p:729-766,
  author={Francisco Gomes and Alexander Michaelides},
  title={{Portfolio Choice With Internal Habit Formation: A Life-Cycle Model With Uninsurable Labor Income Risk}},
  journal={Review of Economic Dynamics},
  year=2003,
  volume={6},
  number={4},
  pages={729-766},
  month={October},
  keywords={Life-cycle asset allocation; Habit formation; Liquidity constraints; Stock market participation cost},
  doi={10.1016/S1094-2025(03)000},
  abstract={Motivated by the success of internal habit formation preferences in explaining asset pricing puzzles, we introduce these preferences in a life-cycle model of consumption and portfolio choice with liquidity constraints, undiversifiable labor income risk and stock-market participation costs. In contrast to the initial motivation, we find that the model is not able to simultaneously match two very important stylized facts: a low stock market participation rate, and moderate equity holdings for those households that do invest in stocks. Habit formation increases wealth accumulation because the intertemporal consumption smoothing motive is stronger. As a result, households start participating in the stock market very early in life, and invest their portfolios almost fully in stocks. Therefore, we conclude that, with respect to its ability to match the empirical evidence on asset allocation behavior, the internal habit formation model is dominated by its time-separable utility counterpart. (Copyright: Elsevier)},
  url={https://ideas.repec.org/a/red/issued/v6y2003i4p729-766.html}
}
'''
```


## **Viscoelastic wave equation implementation on a staggered grid**

This is a first attempt at implementing the viscoelastic wave equation as described in [1]. See also the FDELMODC implementation by Jan Thorbecke [2].

In the following example, a three dimensional toy problem will be introduced consisting of a single Ricker source located at (100, 50, 35) in a 200 m $\times$ 100 m $\times$ 100 *m* domain.

```python
# Required imports:
import numpy as np
import sympy as sp

from devito import *
from examples.seismic.source import RickerSource, TimeAxis
from examples.seismic import ModelViscoelastic, plot_image
```

The model domain is now constructed. It consists of an upper layer of water, 50 m in depth, and a lower rock layer separated by a 4 m thick sediment layer.

```python
# Domain size:
extent = (200., 100., 100.) # 200 x 100 x 100 m domain
h = 1.0 # Desired grid spacing
shape = (int(extent[0]/h+1), int(extent[1]/h+1), int(extent[2]/h+1))

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

# Set up three horizontally separated layers:
vp[:,:,:int(0.5*shape[2])+1] = 1.52
qp[:,:,:int(0.5*shape[2])+1] = 10000.
vs[:,:,:int(0.5*shape[2])+1] = 0.
qs[:,:,:int(0.5*shape[2])+1] = 0.
rho[:,:,:int(0.5*shape[2])+1] = 1.05

vp[:,:,int(0.5*shape[2])+1:int(0.5*shape[2])+1+int(4/h)] = 1.6
qp[:,:,int(0.5*shape[2])+1:int(0.5*shape[2])+1+int(4/h)] = 40.
vs[:,:,int(0.5*shape[2])+1:int(0.5*shape[2])+1+int(4/h)] = 0.4
qs[:,:,int(0.5*shape[2])+1:int(0.5*shape[2])+1+int(4/h)] = 30.
rho[:,:,int(0.5*shape[2])+1:int(0.5*shape[2])+1+int(4/h)] = 1.3

vp[:,:,int(0.5*shape[2])+1+int(4/h):] = 2.2
qp[:,:,int(0.5*shape[2])+1+int(4/h):] = 100.
vs[:,:,int(0.5*shape[2])+1+int(4/h):] = 1.2
qs[:,:,int(0.5*shape[2])+1+int(4/h):] = 70.
rho[:,:,int(0.5*shape[2])+1+int(4/h):] = 2.
```

Now create a Devito vsicoelastic model generating an appropriate computational grid along with absorbing boundary layers:

```python
# Create model
origin = (0, 0, 0)
spacing = (h, h, h)
so = 4 # FD space order (Note that the time order is by default 1).
nbl = 20 # Number of absorbing boundary layers cells
model = ModelViscoelastic(space_order=so, vp=vp, qp=qp, vs=vs, qs=qs,
                          rho=rho, origin=origin, shape=shape, spacing=spacing,
                          nbl=nbl)
```

    Operator `initdamp` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s
    Operator `padfunc` run in 0.01 s

The source frequency is now set along with the required model parameters:

```python
# Source freq. in MHz (note that the source is defined below):
f0 = 0.12

# Thorbecke's parameter notation
l = model.lam
mu = model.mu
ro = model.irho

k = 1.0/(l + 2*mu)
pi = l + 2*mu

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

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

Generate Devito time functions for the velocity, stress and memory variables appearing in the viscoelastic model equations. By default, the initial data of each field will be set to zero.

```python
# PDE fn's:
x, y, z = model.grid.dimensions
damp = model.damp

# Staggered grid setup:

# Velocity:
v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=so)

# Stress:
tau = TensorTimeFunction(name='t', grid=model.grid, space_order=so, time_order=1)

# Memory variable:
r = TensorTimeFunction(name='r', grid=model.grid, space_order=so, time_order=1)

s = model.grid.stepping_dim.spacing # Symbolic representation of the model grid spacing
```

And now the source and PDE's are constructed:

```python
# Source
src = RickerSource(name='src', grid=model.grid, f0=f0, time_range=time_range)
src.coordinates.data[:] = np.array([100., 50., 35.])

# The source injection term
src_xx = src.inject(field=tau[0, 0].forward, expr=src*s)
src_yy = src.inject(field=tau[1, 1].forward, expr=src*s)
src_zz = src.inject(field=tau[2, 2].forward, expr=src*s)

# Particle velocity
u_v = Eq(v.forward, model.damp * (v + s*ro*div(tau)))

# Stress equations:
u_t = Eq(tau.forward,  model.damp *  (r.forward + tau +
                                      s * (l * t_ep / t_s * diag(div(v.forward)) +
                                           mu * t_es / t_s * (grad(v.forward) + grad(v.forward).T))))

# Memory variable equations:
u_r = Eq(r.forward, damp * (r - s / t_s * (r + l * (t_ep/t_s-1) * diag(div(v.forward)) +
                                           mu * (t_es/t_s-1) * (grad(v.forward) + grad(v.forward).T) )))
```

We now create and then run the operator:

```python
# Create the operator:
op = Operator([u_v, u_r, u_t] + src_xx + src_yy + src_zz,
               subs=model.spacing_map)
```

```python
#NBVAL_IGNORE_OUTPUT

# Execute the operator:
op(dt=dt)
```

    Operator `Kernel` run in 12.03 s

Before plotting some results, let us first look at the shape of the data stored in one of our time functions:

```python
v[0].data.shape
```

    (2, 241, 141, 141)

Since our functions are first order in time, the time dimension is of length 2. The spatial extent of the data includes the absorbing boundary layers in each dimension (i.e. each spatial dimension is padded by 20 grid points to the left and to the right).

The total number of instances in time considered is obtained from:

```python
time_range.num
```

    136

Hence 223 time steps were executed. Thus the final time step will be stored in index given by:

```python
np.mod(time_range.num,2)
```

    0

Now, let us plot some 2D slices of the fields `vx` and `szz` at the final time step:

```python
#NBVAL_SKIP

# Mid-points:
mid_x = int(0.5*(v[0].data.shape[1]-1))+1
mid_y = int(0.5*(v[0].data.shape[2]-1))+1

# Plot some selected results:

plot_image(v[0].data[1, :, mid_y, :], cmap="seismic")
plot_image(v[0].data[1, mid_x, :, :], cmap="seismic")

plot_image(tau[2, 2].data[1, :, mid_y, :], cmap="seismic")
plot_image(tau[2, 2].data[1, mid_x, :, :], cmap="seismic")
```

```python
#NBVAL_IGNORE_OUTPUT

assert np.isclose(norm(v[0]), 0.102026, atol=1e-4, rtol=0)
```

# References

[1] Johan O. A. Roberston, *et.al.* (1994). "Viscoelatic finite-difference modeling" GEOPHYSICS, 59(9), 1444-1456.

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


```python
import pickle
import sys
sys.path.insert(0, '../')
from AircraftIden.data_case.GeneralAircraftCase import GeneralAircraftCase, PX4AircraftCase
from AircraftIden import FreqIdenSIMO, TransferFunctionFit,TransferFunctionModel,TransferFunctionParamModel
import numpy as np
import matplotlib.pyplot as plt
import math
import sympy as sp
sp.init_printing()
pkl_name = "/Users/plane/Dropbox/FLYLOG/foam-tail-sitter/cruising/log_34_2018-5-29-freq-res-thrustax.pkl"
thr_freqres = pickle.load(open(pkl_name,"rb"))
```

```python
def process_status_freqres(freqres):
    freq, H, gamma2, gxx, gxy, gyy = freqres.get_freq_iden(0)
    a,b,c,d,tau,s = sp.symbols("a b c d tau s")

    num = d
    den = (b *s + c)
    tfpm = TransferFunctionParamModel(num,den,tau)
    fitter = TransferFunctionFit(freq, H, gamma2, tfpm,nw=20,iter_times=1000,reg = 0)
    tf = fitter.estimate(10,accept_J=10)
    return fitter
#fitter = process_status_freqres(ele_freqres["1m/s"])

```

```python
status_list = ["8m/s","10m/s"]

def process_status(x):
    print("Status process flying status {}".format(x))
    return process_status_freqres(thr_freqres[x])
ret = map(process_status,status_list)
ret = dict(zip(status_list, ret))
```

    Status process flying status 8m/s
    Uknown number 4 [d, c, tau, b]
    Found new better 11.927179919940569
    Status process flying status 10m/s
    Uknown number 4 [d, c, tau, b]
    Found new better 8.344480913890186

```python
plt.rc('figure', figsize=(20.0, 10.0))
def plot_fitter(fitter, status_name):
    plt.figure(status_name)
    fitter.plot(status_name+":")
    trans_str = fitter.latex()
    plt.text(20, 0.8,trans_str, fontsize=20,color='red')
    plt.show()

def poly_latex(poly,cha = "s"):
    ret_str = ""
    ords = len(poly) - 1
    for i in range(ords+1):
        ordn = ords - i
        if ordn == 0:
            ret_str = ret_str + "{:4.2f}".format(poly[i])
        else:
            if (poly[i] == 1):
                ret_str = ret_str + "s^{:d}+".format(ordn)
            else:
                ret_str = ret_str + "{:4.2f} s^{:d}+".format(poly[i],ordn)
    return ret_str

def transfer_func_latex(num,den,tau):
    return r"$\frac{" + poly_latex(num) + "}{" + poly_latex(den) +"}" + "e^{-" + "{:4.3f}".format(tau) + "t}$"

for status in ret:
    plot_fitter(ret[status],status)
```

```python

```


```python
#hide code segments from display
from IPython.display import HTML
HTML('''
The raw code for this notebook is hidden for easier reading.
To toggle code display, click <a href="javascript:code_toggle()">here</a>.''')
```

The raw code for this notebook is hidden for easier reading.
To toggle code display, click here.

# Pen Spring Diffraction Experiment
In this experiment, we illustrate the diffraction of a pen spring (a regular helix) and discuss its implications.

We will first import `numpy` to aid us in our calculations and `matplotlib` to aid us in visualisation.

```python
import numpy as np
import matplotlib.pyplot as plt
import math
%config InlineBackend.figure_format = 'retina'

def sinc(x):
    if (x != 0):
        return np.sin(np.pi * x) / (np. pi * x)
    else:
        return 1
sinc = np.vectorize(sinc)
```

# 1. Fill in: Experiment Variables
Enter the values here to match your experimental apparatus.

1. `slitWidth` is the gap in between each individual coil of the spring. From now on, it will be referred to as the *pitch* of the spring.
2. `wavelength` is the wavelength of the laser that is emitted by the laser pointer. It is usually marked near the emitter, on a warning label.
3. `distance` is the distance between the center of the spring and the surface on which the pattern is projection upon.

Units of the values below are given in **micrometers** ($10^{-6}$ m, also refered to as the **micron**).

```python
'''
slitWidth  = float(input("slitWidth: ")) #2500          # 2.5mm
wavelength = float(input("wavelength: ")) #0.650         # 650nm
distance   = float(eval(input("distance: "))) #10.3 * 10**6  # 10.3m
'''
slitWidth  = 2500          # 2.5mm
wavelength = 0.650         # 650nm
distance   = 10.3 * 10**6  # 10.3m
```

# 2. Fill in: Experiment Results

Measure the distance from the center to the $n-$th maxima. Choose any reasonable number!

Change the variable here to match your result in **mm**. It will be plotted on the visualisation later on.

```python
#result_mm = float(input("result: "))
result_mm = 25.5
```

# Visualising a Single-slit Diffraction
We will first set up some parameters, and then proceed to plot the intensity over distance from the center of the pattern.

```python
amplitude = 5 # the power of the laser (mW) as marked on the laser point. The accuracy of this value is trivial.
x = np.arange(-50000, 50000, 10) # numerical bounds and resolution of the graph (micron)
F = sinc(slitWidth * x / wavelength / distance)
intensity = amplitude / (wavelength * distance) * (slitWidth * F)**2
```

### Plotting the graph
See if your result (marked by blue line) indeed corresponds to the maxima from which you measured earlier on

```python
plt.figure(1, (12, 6))
plt.plot(x, intensity, linewidth=1, color='r') # red indicates the intensity
result = result_mm * 1000
plt.axvline(x = result, color='b', label="Obtained result: "+str(result_mm)+"mm") # blue indicated your result indicated earlier
plt.xlim((0, 50000))
plt.ylim((0, 0.06))
plt.xlabel(r'Position ($\mu$m)m)')
plt.ylabel(r'Intensity (mW)')
plt.legend(loc='lower right')
plt.title("Single-slit Diffraction (intensity / distance)")
plt.show()
```

# Angle between the Lines
Now, we will attempt to calculate the angle between the lines that form the X-shape. It corresponds to the *helix angle* on the spring.

The laser will illuminate one pitch (complete curl) of the spring.

**Babinet’s principle** states that the diffraction pattern of an obstacle is the same as the diffraction pattern of an aperture of the same shape. Hence, we will expect to see *two* single slit diffraction patterns angled to each other.

When viewed from the side, the spring (or any regular helix) resembles a sine wave. We also will make use of this property later on.

```python
plt.figure(1, (3, 4))
t = np.arange(0, 2.5, 0.04)
s = np.sin(2*np.pi*t)
plt.plot(s, t, color='k', linewidth=2)
x_p = np.linspace(-0.9, 0.9, 3)
line1,  = plt.plot(x_p, 0.17*x_p + 1, '--', linewidth=4, color='r', label='Illuminated Section')
line2,  = plt.plot(x_p, -0.17*x_p + 0.98+np.pi/6, '--', linewidth=4, color='r')
plt.axis('off')
plt.legend(loc='lower right')
plt.title("Sideways projection of the Spring")
plt.show()
```

# 3. Fill In: Parameters of the Spring
Enter the dimensions of the spring below, measured with a pair of vernier calipers. The units here are given in **mm**.

```python
'''
pitch_mm  = float(input("pitch_mm: "))
radius_mm = float(input("radius_mm: "))
'''
pitch_mm  = 2.5   # p
radius_mm = 2.05  # R
```

We orientate the spring in the $y-z$ plane to make it easier to work with.

The code below generates a graphical representation of your spring projected onto the $y-z$ plane.

```python
plt.figure(2, (12, 8))
z_q = np.arange(-2*pitch_mm, 2*pitch_mm, 0.05)
y_q = radius_mm * np.sin(2*np.pi * z_q/pitch_mm)
plt.plot(z_q, y_q, color = 'k', linewidth=4 , label='Spring')

y_1 = radius_mm * (2*np.pi * z_q/pitch_mm - 2*np.pi)
y_2 = radius_mm * (2*np.pi * z_q/pitch_mm)
y_3 = radius_mm * (2*np.pi * z_q/pitch_mm + 2*np.pi)

alpha = np.pi/2 - math.atan(radius_mm * 2*np.pi / pitch_mm)

plt.plot(z_q, y_1, "--", color='r', label = r'$y=\frac{2\pi R}{p}z+2n\pi$')
plt.plot(z_q, y_2, "--", color='r')
plt.plot(z_q, y_3, "--", color='r')

plt.text(0.1, 2, r'$\alpha$', fontsize=20)

d = pitch_mm*math.cos(alpha)
x_d2 = pitch_mm*(math.cos(alpha)**2)

x_d = np.linspace(0, x_d2, 3)
plt.plot(x_d, -1*(pitch_mm/(2*np.pi*radius_mm)) * x_d, "--", color='b')
plt.text(0.1, -radius_mm/3, "$d$ = " + str('%.1f'%d) + "mm", color="b", fontsize=14)

x_p1 = pitch_mm/4
x_p2 = x_p1 + pitch_mm
x_pitch = np.linspace(x_p1, x_p2, 3)
plt.plot(x_pitch, radius_mm + x_pitch*0, "--", color='g')
plt.text((x_p1+x_p2)/2-0.6, radius_mm + 0.1, "$p$ = " + str(pitch_mm) + "mm", color="g", fontsize=14)

plt.xlabel('$z$ (mm)')
plt.axvline(x=0, color='k', alpha=0.5)
plt.ylabel('$y$ (mm)')
plt.axhline(y=0, color='k', alpha=0.5)
plt.axis('equal')
plt.xlim((-2*pitch_mm, 2*pitch_mm))
plt.ylim((-radius_mm*1.2, radius_mm*1.2))
plt.legend(loc='lower right')
ax = plt.gca()
ax.set_xticks(np.arange(-2*pitch_mm, 2*pitch_mm, 1))
ax.set_yticks(np.arange(-radius_mm*2, radius_mm*2, 1))
plt.grid(True, linestyle='--', linewidth=1, alpha=0.5)
plt.title(r"Projection of Spring [$p=$"+str(pitch_mm)+r"mm, $R=$"+str(radius_mm)+r"mm] in the $y-z$ plane")
plt.show()
```

Perform some symbolic differentiation on the equation of the spring:

```python
from sympy import *
init_printing()
R, p, z = symbols('R p z')
eqn_y = R * sin(2 * pi * z / p)
eqn_y.simplify()
```

```python
n = symbols('n', integer=True)
deriv_y = diff(eqn_y, z)
deriv_y.simplify()
```

Substitute $z=n \times p$:

```python
deriv_y = deriv_y.subs(z, n*p) #z = n*p
deriv_y.simplify()
```

The angle $\alpha$ is equal to the following expression:

```python
eqn_a = -1 * atan(deriv_y) + pi/2
eqn_a.simplify()
```

```python
deriv_y = deriv_y.subs(p, pitch_mm * 10**-3)
deriv_y = deriv_y.subs(R, radius_mm * 10**-3)
a = -1 * math.atan(deriv_y) + math.pi/2
print(r'The angle(α) between the lines is '
      + str('%.3f'%(math.degrees(a)) + " degrees (to 3dp)."))
```

    The angle(α) between the lines is 10.984 degrees (to 3dp).

# Visualising the Actual Diffraction Pattern
Now that we have the angle and the intensity, we can attempt to visualise the pattern seen.

```python
angle = math.pi/2 - a

plt.figure(3, (12, 12))

# transform the intensity line by rotating it about the center
x_tf = ( x / math.sqrt((math.tan(angle))**2 + 1) )
y_tf = x_tf * math.tan(angle)
plt.scatter(x_tf, y_tf, color='r', s=intensity, alpha=0.05) # plot the first line

y_tf = x_tf * math.tan(-angle)
plt.scatter(x_tf, y_tf, color='r', s=intensity, alpha=0.1) # plot the second line

# plot the circle which marks our result from earlier (gray)
ax = plt.gca()
circle_result = plt.Circle((0, 0), result, color='b', fill=False, alpha=0.2)
ax.add_artist(circle_result)

plt.xlabel("$x$")
plt.ylabel("$y$")
plt.axis('equal')
plt.xlim((-30000, 30000))
plt.ylim((-30000, 30000))
plt.show()
```

#### Zoomed in version
This is a close-up look at the first quadrant of the above figure.

It allows us to properly judge the maxima that might not be obvious above.

```python
plt.figure(4, (12, 12))

x_tf = ( x / math.sqrt((math.tan(angle))**2 + 1) )
y_tf = x_tf * math.tan(angle)
plt.scatter(x_tf, y_tf, color='r', s=intensity*100, alpha=0.05)

# circle which marks the result (blue)
ax = plt.gca()
circle_result = plt.Circle((0, 0), result, color='b', fill=False)
ax.add_artist(circle_result)

plt.xlabel("$x$")
plt.ylabel("$y$")
plt.axis('equal')
plt.xlim((0, 1.1*result))
plt.ylim((0, 30000))
plt.grid(True, linestyle='--', linewidth=1, alpha=0.5)
plt.show()
```

# Significance of this Result

**Photograph 51** is the nickname given to an X-ray diffraction image of DNA taken by Raymond Gosling in May 1952, working as a PhD student under the supervision of Rosalind Franklin.

Does it not resemble the result you have gotten from the pen spring?

Since objects of similar shape produce similar diffraction patterns (Babinet's principle), this result was critical evidence in identifying the structure of DNA as a regular helix.

```python
import os
print(os.popen("cowsay 'Can MatLab do this?'").read())
```

     _____________________
    < Can MatLab do this? >
     ---------------------
            \   ^__^
             \  (oo)\_______
                (__)\       )\/\
                    ||----w |
                    ||     ||