 # 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}
}
'''
```

    '\n@Article{RePEc:red:issued:v:6:y:2003:i:4:p:729-766,\n  author={Francisco Gomes and Alexander Michaelides},\n  title={{Portfolio Choice With Internal Habit Formation: A Life-Cycle Model With Uninsurable Labor Income Risk}},\n  journal={Review of Economic Dynamics},\n  year=2003,\n  volume={6},\n  number={4},\n  pages={729-766},\n  month={October},\n  keywords={Life-cycle asset allocation; Habit formation; Liquidity constraints; Stock market participation cost},\n  doi={10.1016/S1094-2025(03)000},\n  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)},\n  url={https://ideas.repec.org/a/red/issued/v6y2003i4p729-766.html}\n}\n'


```python
from controleII import digital
from lcapy.discretetime import z, n
from lcapy import s,exp
from control import tf
```

```python
# Questão 1
Xz = (10 * z + 5) / ((z-1) * (z-0.2))
teste = digital(Xz)
display(teste.izt())
print(teste.calcn())
```

```python
# Questão 2
Xz = 1 + 2*z**-1 + 3*z**-2 + 4 *z **-3
teste = digital(Xz)
display(teste.izt())
print(teste.calcn())
```

```python
# Questão 4

Xz = (z**2 + z + 2) / ((z-1)* (z**2 -z + 1))
Q4 = digital(Xz=Xz)
display(Q4.izt())
Q4.calcn()
```

```python
# Questão 5
Xz = (2*z**3 + z) / ((z-2)**2 * (z-1))
Q5 = digital(Xz = Xz)
display(Q5.izt())
```

```python
# Questão 6
Xz1 = 9 / (z * (1-2*z**-1)**2)
Xz2 = z / (z - 2)
Xz3 = 3*z / (z - 1)
Xz = Xz1 + Xz2 + Xz3
teste = digital(Xz)
display(teste.izt())
```

```python
# Questão 7

Xz = (z+2) / (z**2 * (z-2))
Q7 = digital(Xz = Xz)
display(Q7.izt())
```

```python
# Questão 13:
gs = 1 / (s * (s + 1))
gsc = tf([1], [1, 1])
teste = digital(Gs = gs, Gsc = gsc, tau=1)
display(teste.s2t())
display(teste.s2z())
```

```python
# questão 14

gss = 4 / (s * (s+2))
gsc = tf(4, [1, 2])
tau = .1
Q14 = digital(Gs=gss, Gsc=gsc, tau = .1)
display(Q14.s2t())
display(Q14.s2z())
fz, time, mag = Q14.Fz(20)
display(fz, time, mag)
```

```python
# questão 16

gss = 10 / (s + 10*s**2)
gsc = tf([10], [10, 1])
tau = 2
Q16 = digital(Gs=gss, Gsc=gsc, tau = 2)
display(Q16.s2z())
display(Q16.s2t())
fz, time, mag = Q16.Fz(14, fb = 0.05, plot=True)
display(fz, mag)

```

```python
g1s = 2 * (s - 6.93) / (s - 1) * (1 / (s * (s+1)))
gsc =  tf(1, [1, 1]) #* 2* tf([1, -6.93], [1, -1])
tau = .1
Q17 = digital(Gs=g1s, Gsc=gsc, tau = tau)
display(Q17.s2z())
fz, time, mag = Q17.Fz(30, plot=True, save=True, title='Resposta ao degrau unitário sem o controlador')
display(fz)

```

```python
Xz = (0.130976*z - 0.26969) / (z**2 - 1.879029*z + 0.7303) * (z / (z-1))
C = N(Xz.partfrac())
C
```

```python
from sympy import solve, var, N
var('u')
a =(u**2 - 2.01*u + 1) * u/(u-1)
solve(u, a)
```

```python
from sympy import var
from sympy import solve

k = var('k')
W = var('W')
t = (2 + W) / (2 - W)

Xz = k/(z-1) - ((1-exp(-1))/(z- exp(-1)))*k
#Xz = Xz.subs(z, t)
F = Xz/(Xz + 1)
G = F.simplify().as_numer_denom()[1]
G.subs(k, 25)
```


```python
import sys
import getopt
import os
#from sets import Set
import json
import pprint
import itertools
from graphviz import Digraph
from math import log
import src.searchtree_v2 as st
import src.utils as utils
import csv
from pathlib import Path

```

```python
#compose the logic formula describing the properties of the operators
from sympy import *

#helper functions
def getPreSATVariable(op,variable,value):
    var = op+"preVar"+str(variable)+"Val"+str(value)
    sym = symbols(var)
    return sym

def getEffSATVariable(op,variable,value):
    var = op+"effVar"+str(variable)+"Val"+str(value)
    sym = symbols(var)
    return sym

def opInitApplicableInState(opHash):
    formula = True
    for var in range(0,agentPrivateVariables):
        sym = getPreSATVariable(opHash,var,0)
        formula = (formula & sym)
    return formula

def opPrivatelyIndependent(opHash,numOfStates):
    formula = False
    for var in range(0,agentPrivateVariables):
        subformula = True
        for val in range(0,agentDomainSize):
            sym = getPreSATVariable(opHash,var,val)
            subformula = (subformula & sym)
        formula = formula | subformula
    return formula

def opPrivatelyDependent(opHash,numOfStates):
    formula = False
    for var in range(0,agentPrivateVariables):
        subformula1 = True
        subformula2 = True
        for val in range(0,agentDomainSize):
            sym = getPreSATVariable(opHash,var,val)
            subformula1 = subformula1 | sym
            subformula2 = subformula2 | ~sym
        formula = formula | (subformula1 & subformula2)
    return formula

def opPrivatelyNondeterministic(opHash,numOfStates):
    formula = False
    for var in range(0,agentPrivateVariables):
        subformula = False
        for val1 in range(0,agentDomainSize):
            for val2 in range(0,agentDomainSize):
                if val1 != val2:
                    sym1 = getEffSATVariable(opHash,var,val1)
                    sym2 = getEffSATVariable(opHash,var,val2)
                    subformula = subformula | (sym1 & sym2)
        formula = formula | (subformula)
    return formula

```

```python
def computeFixedTS(k,l):

    sum = 0
    for j in range(k+1):
        sum += (-1)**j * binomial(k,j) * (2**(k-j) - 1)**l

#     print("compute all TS between "+str(k)+" pre and "+str(l)+" eff = " + str(sum))

    return sum

def computeAllTS(pfree,efree,ptrue,etrue):
    total = 0

    pmin = 0
    if ptrue==0:
        pmin=1

    emin = 0
    if etrue==0:
        emin=1

    for k in range(pmin,pfree+1):
        for l in range(emin,efree+1):
            if k+ptrue>0 and l+etrue >0:
                fixed = binomial(pfree,k) * binomial(efree,l) * computeFixedTS(k+ptrue,l+etrue)
                total += fixed

    return total

def evaluateModels(models):
    count = 0;
    maxTotalTS = 0

    for model in models:
#         print(model)
        count += 1
        opHash = 'op'

        opTS = 1
        subForEmptyTS = 1

        if model == True or model == False:
            return 1

        for var in range(0,agentPrivateVariables):

            freePre = agentDomainSize
            freeEff = agentDomainSize

            posPre = 0
            posEff = 0

            for val in range(0,agentDomainSize):
                preVar = getPreSATVariable(opHash,var,val)
                effVar = getEffSATVariable(opHash,var,val)

                if preVar in model:
                    freePre -= 1
                    if model[preVar] == True:
                        posPre += 1

                if effVar in model:
                    freeEff -=1
                    if model[effVar] == True:
                        posEff += 1

#                 print(opHash+",var="+str(var)+": pre:#"+str(freePre)+",+"+str(posPre))
#                 print(opHash+",var="+str(var)+": eff:#"+str(freeEff)+",+"+str(posEff))

            varTS = computeAllTS(freePre,freeEff,posPre,posEff)

            opTS = opTS * varTS

        totalTS = opTS

        if totalTS > maxTotalTS:
            maxTotalTS = totalTS

    print("Evaluated " + str(count) + " models")
    return maxTotalTS

#test:
print(computeAllTS(2,2,0,0))
```

    15

```python
#determine RHS based on combinations of properties

#check cache:
cacheCSVFile = "ts_cache.csv"
outCSV = Path(cacheCSVFile)
exists = outCSV.is_file()
found = False

for agentPrivateVariables in range (1,5):
    for agentDomainSize in range(1,5):
        print("Compute: " + str(agentPrivateVariables)+","+str(agentDomainSize))

        found = False
        if exists:
            with open("ts_cache.csv") as csvfile:
                reader = csv.DictReader(csvfile)
                for row in reader:
                    if row['agentDomainSize']==str(agentDomainSize) and  row['agentPrivateVariables']==str(agentPrivateVariables):
                        print("Reading from cache file!")
                        print(row)
                        found = True

        if found:
            print("Found!")
            continue

        print("ts")
        ts = (2**(agentDomainSize**2) - 1)**agentPrivateVariables
        print(ts)
        rhs_t = log(ts,2).evalf()

        print("ia")
        f_ia = opInitApplicableInState('op')
        print(f_ia)
        models_ia = satisfiable(f_ia,all_models=True)
        ts_ia = evaluateModels(models_ia)
        print(ts_ia)
        rhs_ia = log(ts_ia,2).evalf()

        print("nia")
        f_nia = ~opInitApplicableInState('op')
        print(f_nia)
        models_nia = satisfiable(f_nia,all_models=True)
        ts_nia = evaluateModels(models_nia)
        print(ts_nia)
        rhs_nia = log(ts_nia,2).evalf()

        print("pi")
        f_pi = opPrivatelyIndependent('op',1)
        print(f_pi)
        models_pi = satisfiable(f_pi,all_models=True)
        ts_pi = evaluateModels(models_pi)
        print(ts_pi)
        rhs_pi = log(ts_pi,2).evalf()

        print("pd")
        f_pd = opPrivatelyDependent('op',1)
        print(f_pd)
        models_pd = satisfiable(f_pd,all_models=True)
        ts_pd = evaluateModels(models_pd)
        print(ts_pd)
        rhs_pd = log(ts_pd,2).evalf()

        print("pn")
        f_pn = opPrivatelyNondeterministic('op',1)
        print(f_pn)
        models_pn = satisfiable(f_pn,all_models=True)
        ts_pn = evaluateModels(models_pn)
        print(ts_pn)
        rhs_pn = log(ts_pn,2).evalf()

        print("ia_pi")
        f_ia_pi = opInitApplicableInState('op') & opPrivatelyIndependent('op',1)
        print(f_ia_pi)
        models_ia_pi = satisfiable(f_ia_pi,all_models=True)
        ts_ia_pi = evaluateModels(models_ia_pi)
        print(ts_ia_pi)
        rhs_ia_pi = log(ts_ia_pi,2).evalf()

        print("nia_pi")
        f_nia_pi = ~opInitApplicableInState('op') & opPrivatelyIndependent('op',1)
        print(f_nia_pi)
        models_nia_pi = satisfiable(f_nia_pi,all_models=True)
        ts_nia_pi = evaluateModels(models_nia_pi)
        print(ts_nia_pi)
        rhs_nia_pi = log(ts_nia_pi,2).evalf()

        print("ia_pd")
        f_ia_pd = opInitApplicableInState('op') & opPrivatelyDependent('op',1)
        print(f_ia_pd)
        models_ia_pd = satisfiable(f_ia_pd,all_models=True)
        ts_ia_pd = evaluateModels(models_ia_pd)
        print(ts_ia_pd)
        rhs_ia_pd = log(ts_ia_pd,2).evalf()

        print("nia_pd")
        f_nia_pd = ~opInitApplicableInState('op') & opPrivatelyDependent('op',1)
        print(f_nia_pd)
        models_nia_pd = satisfiable(f_nia_pd,all_models=True)
        ts_nia_pd = evaluateModels(models_nia_pd)
        print(ts_nia_pd)
        rhs_nia_pd = log(ts_nia_pd,2).evalf()

        print("ia_pn")
        f_ia_pn = opInitApplicableInState('op') & opPrivatelyNondeterministic('op',1)
        print(f_ia_pn)
        models_ia_pn = satisfiable(f_ia_pn,all_models=True)
        ts_ia_pn = evaluateModels(models_ia_pn)
        print(ts_ia_pn)
        rhs_ia_pn = log(ts_ia_pn,2).evalf()

        print("nia_pn")
        f_nia_pn = ~opInitApplicableInState('op') & opPrivatelyNondeterministic('op',1)
        print(f_nia_pn)
        models_nia_pn = satisfiable(f_nia_pn,all_models=True)
        ts_nia_pn = evaluateModels(models_nia_pn)
        print(ts_nia_pn)
        rhs_nia_pn = log(ts_nia_pn,2).evalf()

        print("pi_pn")
        f_pi_pn = opPrivatelyIndependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_pi_pn)
        models_pi_pn = satisfiable(f_pi_pn,all_models=True)
        ts_pi_pn = evaluateModels(models_pi_pn)
        print(ts_pi_pn)
        rhs_pi_pn = log(ts_pi_pn,2).evalf()

        print("pd_pn")
        f_pd_pn = opPrivatelyDependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_pd_pn)
        models_pd_pn = satisfiable(f_pd_pn,all_models=True)
        ts_pd_pn = evaluateModels(models_pd_pn)
        print(ts_pd_pn)
        rhs_pd_pn = log(ts_pd_pn,2).evalf()

        print("ia_pi_pn")
        f_ia_pi_pn = opInitApplicableInState('op') & opPrivatelyIndependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_ia_pi_pn)
        models_ia_pi_pn = satisfiable(f_ia_pi_pn,all_models=True)
        ts_ia_pi_pn = evaluateModels(models_ia_pi_pn)
        print(ts_ia_pi_pn)
        rhs_ia_pi_pn = log(ts_ia_pi_pn,2).evalf()

        print("nia_pi_pn")
        f_nia_pi_pn = ~opInitApplicableInState('op') & opPrivatelyIndependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_nia_pi_pn)
        models_nia_pi_pn = satisfiable(f_nia_pi_pn,all_models=True)
        ts_nia_pi_pn = evaluateModels(models_nia_pi_pn)
        print(ts_nia_pi_pn)
        rhs_nia_pi_pn = log(ts_nia_pi_pn,2).evalf()

        print("ia_pd_pn")
        f_ia_pd_pn = opInitApplicableInState('op') & opPrivatelyDependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_ia_pd_pn)
        models_ia_pd_pn = satisfiable(f_ia_pd_pn,all_models=True)
        ts_ia_pd_pn = evaluateModels(models_ia_pd_pn)
        print(ts_ia_pd_pn)
        rhs_ia_pd_pn = log(ts_ia_pd_pn,2).evalf()

        print("nia_pd_pn")
        f_nia_pd_pn = ~opInitApplicableInState('op') & opPrivatelyDependent('op',1) & opPrivatelyNondeterministic('op',1)
        print(f_nia_pd_pn)
        models_nia_pd_pn = satisfiable(f_nia_pd_pn,all_models=True)
        ts_nia_pd_pn = evaluateModels(models_nia_pd_pn)
        print(ts_nia_pd_pn)
        rhs_nia_pd_pn = log(ts_nia_pd_pn,2).evalf()

        columns=[
            'agentPrivateVariables',
            'agentDomainSize',
            'ts',
            'ia',
            'nia',
            'pi',
            'pd',
            'pn',
            'ia_pi',
            'ia_pd',
            'ia_pn',
            'nia_pi',
            'nia_pd',
            'nia_pn',
            'pi_pn',
            'pd_pn',
            'ia_pi_pn',
            'ia_pd_pn',
            'nia_pi_pn',
            'nia_pd_pn'
        ]

        outCSV = Path(cacheCSVFile)
        exists = outCSV.is_file()

        row = [
            agentPrivateVariables,
            agentDomainSize,
            rhs_t,
            rhs_ia,
            rhs_nia,
            rhs_pi,
            rhs_pd,
            rhs_pn,
            rhs_ia_pi,
            rhs_ia_pd,
            rhs_ia_pn,
            rhs_nia_pi,
            rhs_nia_pd,
            rhs_nia_pn,
            rhs_pi_pn,
            rhs_pd_pn,
            rhs_ia_pi_pn,
            rhs_ia_pd_pn,
            rhs_nia_pi_pn,
            rhs_nia_pd_pn
        ]

        with open(cacheCSVFile, 'a') as f:
            writer = csv.writer(f)
            if not exists:
                writer.writerow(columns)
            writer.writerow(row)

```

```python

```


```python
from Exercise import Exercise, MarkdownBlock
from config import URL, TOKEN
import json

import numpy as np
import sympy as sp

Exercise.URL = URL
Exercise.TOKEN = TOKEN

MarkdownBlock.TEX_INLINE_WRAP_LEFT = "\("
MarkdownBlock.TEX_INLINE_WRAP_RIGHT = "\)"
```

```python
from sympy import Rational, Symbol, latex, UnevaluatedExpr

u = lambda x : UnevaluatedExpr(x)

# Helper functions
def explain_add(a, b):
    assert(np.shape(a) == np.shape(b))
    rows, columns = np.shape(a)
    return sp.Matrix([[Symbol(f"({latex(u(a[i,j]))} + {latex(u(b[i,j]))})") for j in range(columns)] for i in range(rows)])

def symbolic_matrix(character, rows, columns):
    return sp.Matrix([[Symbol(f"{{{character}}}_{{{i+1}, {j+1}}}") for j in range(columns)] for i in range(rows)])
```

```python
a = np.arange(4)
b = -1 * np.arange(1, 5)
a = sp.Matrix(a)
b = sp.Matrix(b)

s = "What is $@a + @b?$"

params = {}
params["a"] = a
params["b"] = b
e = Exercise(MarkdownBlock(s, params))

s1 = "Yes, $@a + @b = @c = @d$!"
params["c"] = explain_add(a,b)
params["d"] = a + b
e.add_answer(a + b, True, MarkdownBlock(s1, params))

params_1 = {}
params_1["x"] = symbolic_matrix("a", 4,1)
params_1["y"] = symbolic_matrix("b", 4,1)
params_1["z"] = explain_add(params_1["x"], params_1["y"])
default_feedback = "Remember the definition of matrix addition: $@x + @y = @z$"
e.add_default_feedback(MarkdownBlock(default_feedback, params_1))

e.write()
```

```python

```

```python

```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, G.F. Forsyth, C.D. Cooper.

# Spreading out

Welcome back! This is the third lesson of the course [Module 4](https://github.com/numerical-mooc/numerical-mooc/tree/master/lessons/04_spreadout), _Spreading out: parabolic PDEs,_ where we study the numerical solution of diffusion problems.

In the first two notebooks, we looked at the 1D heat equation, and solved it numerically using [*explicit*](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/01_Heat_Equation_1D_Explicit.ipynb) and [*implicit*](https://github.com/numerical-mooc/numerical-mooc/blob/master/lessons/04_spreadout/02_Heat_Equation_1D_Implicit.ipynb) schemes. We learned that implicit schemes are unconditionally stable, and we are free to choose any time step. —Wait: _any time step?_ Remember, we still want to capture the physics of the problem accurately. So although stability concerns do not limit the time step, it still hast to be small enough to satisfy any accuracy concerns.

We are now ready to graduate to two dimensions! In the remaining lessons of this course module, we will study the 2D heat equation and reaction-diffusion equation. Like before, we start with explicit methods (this lesson) and then move to implicit methods (next lesson). Let's get started.

## 2D Heat conduction

The equation of heat conduction in 2D is:

\begin{equation}
\rho c_p \frac{\partial T}{\partial t} = \frac{\partial}{\partial x} \left( \kappa_x \frac{\partial T}{\partial x} \right) + \frac{\partial}{\partial y} \left(\kappa_y \frac{\partial T}{\partial y} \right)
\end{equation}

where $\rho$ is the density, $c_p$ is the heat capacity and $\kappa$ is the thermal conductivity.

If the thermal conductivity $\kappa$ is constant, then we can take it outside of the spatial derivative and the equation simplifies to:

\begin{equation}
\frac{\partial T}{\partial t} = \alpha \left(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right)
\end{equation}

where $\alpha = \frac{\kappa}{\rho c_p}$ is the thermal diffusivity.  The thermal diffusivity describes the ability of a material to conduct heat vs. storing it.

Does that equation have a familiar look to it?  That's because it's the same as the diffusion equation.  There's a reason that $\alpha$ is called the thermal *diffusivity*! We're going to set up a interesting problem where 2D heat conduction is important, and set about to solve it with explicit finite-difference methods.

### Problem statement

Removing heat out of micro-chips is a big problem in the computer industry. We are at a point in technology where computers can't run much faster because the chips might start failing due to the high temperature. This is a big deal! Let's study the problem more closely.

We want to understand how heat is dissipated from the chip with a very simplified model. Say we consider the chip as a 2D plate of size $1{\rm cm}\times 1{\rm cm}$, made of Silicon: $\kappa = 159{\rm W/m C}$, $c_p = 0.712\cdot 10^3 {\rm J/kg C}$, $\rho = 2329{\rm kg/m}^3$, and diffusivity $\alpha \approx 10^{-4}{\rm m}^2{/rm s}$. Silicon melts at $1414{\rm C}$, but chips should of course operate at much smaller temperatures. The maximum temperature allowed depends on the processor make and model; in many cases, the maximum temperature is somewhere between $60{\rm C}$ and $\sim70{\rm C}$, but better CPUs are recommended to operate at a [maximum of $80{\rm C}$](http://www.pugetsystems.com/blog/2009/02/26/intel-core-i7-temperatures/) (like the Intel Core i7, for example).

We're going to set up a somewhat artificial problem, just to demonstrate an interesting numerical solution. Say the chip is in a position where on two edges (top and right) it is in contact with insulating material.  On the other two edges the chip is touching other components that have a constant temperature of $T=100{\rm C}$ when the machine is operating. Initially, the chip is at room temperature $(20{\rm C})$. *How long does it take for the center of the chip to reach $70{\rm C}$?*

####Figure 1: Simplified microchip problem setup.

Let's use what we have learned to tackle this problem!

## 2D Finite differences

Everything you learned about finite-difference schemes in [Notebook 1 of Module 2](http://nbviewer.ipython.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/02_spacetime/02_01_1DConvection.ipynb) still applies, but now there are two spatial dimensions.  We will need to build a 2D grid of discrete points to compute the solution on.

We will use a 2D Cartesian grid: one that consists of two families of (grid) lines parallel to the two spatial directions. Two lines (of different families) intersect on one and only one grid node (this is called a _structured_ grid). In the $x$ direction, the discretization uses $i=0, \cdots N_x$ lines, and in the $y$ direction we have $j=0, \cdots N_y$ lines. A given node on the grid will now have two spatial coordinates, and we need two indices: for the two lines that intersect at that node.  For example, the middle point in the figure below would be $T_{i,j}$.

#### Figure 2. Nodal coordinates in 2 dimensions

### Explicit scheme in 2D

Recall from above that the 2D heat equation is

$$\frac{\partial T}{\partial t} = \alpha \left(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right)$$

Let's write this out discretized using forward difference in time, and central difference in space,  using an explicit scheme. You should be able write this out yourself, without looking—if you need to look, it means you still need to write more difference equations by your own hand!

\begin{equation}
\frac{T^{n+1}_{i,j} - T^n_{i,j}}{\Delta t} = \alpha \left( \frac{T^n_{i+1, j} - 2T^n_{i,j} + T^n_{i-1,j}}{\Delta x^2} + \frac{T^n_{i, j+1} - 2T^n_{i,j} + T^n_{i,j-1}}{\Delta y^2}\right)
\end{equation}

Rearranging the equation to solve for the value at the next time step, $T^{n+1}_{i,j}$, yields

\begin{equation}
T^{n+1}_{i,j}= T^n_{i,j} + \alpha \left( \frac{\Delta t}{\Delta x^2} (T^n_{i+1, j} - 2T^n_{i,j} + T^n_{i-1,j}) + \frac{\Delta t}{\Delta y^2} (T^n_{i, j+1} - 2T^n_{i,j} + T^n_{i,j-1})\right)
\end{equation}

That's a little messier than 1D, but still recognizable.

Up until now, we've used stencils to help visualize how a scheme will advance the solution for one timestep.  Stencils in 2D are a little harder to draw, but hopefully the figure below will guide your understanding of this method: we are using five grid points at time step $n$ to obtain the solution on one point at time step $n+1$.

####Figure 3: 2D Explicit Stencil

Similar to all of the 1D explicit methods we've used, the solution at $T^{n+1}_{i,j}$ is updated using only known values from the current solution at time $n$. This is straightforward to implement in code, but will be subject to stability limitations on the time step that you can choose. We'll study an implicit method in the next lesson.

### Boundary Conditions

Whenever we reach a point that interacts with the boundary, we apply the boundary condition. As in the previous notebook, if the boundary has Dirichlet conditions, we simply impose the prescribed temperature at that point. If the boundary has Neumann conditions, we approximate them with a finite-difference scheme.

Remember, Neumann boundary conditions prescribe the derivative in the normal direction. For example, in the problem described above, we have $\frac{\partial T}{\partial y} = q_y$ in the top boundary and $\frac{\partial T}{\partial x} = q_x$ in the right boundary, with $q_y = q_x = 0$ (insulation).

Thus, at every time step, we need to enforce

\begin{equation}
T_{i,end} = q_y\cdot\Delta y + T_{i,end-1}
\end{equation}

and

\begin{equation}
T_{end,j} = q_x\cdot\Delta x + T_{end-1,j}
\end{equation}

Write the finite-difference discretization of the boundary conditions yourself, and confirm that you can get the expressions above.

### Stability

Before doing any coding, let's revisit stability constraints. We saw in the first notebook of this series that the 1D explicit discretization of the diffusion equation was stable as long as $\alpha \frac{\Delta t}{(\Delta x)^2} \leq \frac{1}{2}$. In 2D, this constraint is even tighter, as we need to add them in both directions:

\begin{equation}
\alpha \frac{\Delta t}{(\Delta x)^2} + \alpha \frac{\Delta t}{(\Delta y)^2} < \frac{1}{2}.
\end{equation}

Say that the mesh has the same spacing in $x$ and $y$, $\Delta x = \Delta y = \delta$. In that case, the stability condition is:

\begin{equation}
\alpha \frac{\Delta t}{\delta^2} < \frac{1}{4}
\end{equation}

## Code implementation

### Array storage

The physical problem has two dimensions, so we also store the temperatures in two dimensions: in a 2D array.

We chose to store it with the $y$ coordinates corresponding to the rows of the array and $x$ coordinates varying with the columns (this is just a code design decision!). If we are consistent with the stencil formula (with $x$ corresponding to index $i$ and $y$ to index $j$), then $T_{i,j}$ will be stored in array format as `T[j,i]`.

This might be a little confusing as most of us are used to writing coordinates in the format $(x,y)$, but our preference is to have the data stored so that it matches the physical orientation of the problem. Then, when we make a plot of the solution, the visualization will make sense to us, with respect to the geometry of our set-up. That's just nicer than to have the plot rotated!

####Figure 4: Row-column data storage

As you can see on Figure 4 above, if we want to access the value $18$ we would write those coordinates as $(x_2, y_3)$.  You can also see that its location is the 3rd row, 2nd column, so its array address would be `T[3,2]`.

Again, this is a design decision.  However you can choose to manipulate and store your data however you like; just remember to be consistent!

### Code time!

Now, to some coding! First, we have a little function that will advance the solution in time with a forward-time, centered-space scheme, and will monitor the center of the plate to tell us when it reaches $70{\rm C}$. Let's start by setting up our Python compute environment.

```python
import numpy
from matplotlib import pyplot
%matplotlib inline
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
```

```python
def ftcs(T, nt, alpha, dt, dx, dy):

    j_mid = (numpy.shape(T)[0])/2
    i_mid = (numpy.shape(T)[1])/2

    for n in range(nt):
        Tn = T.copy()
        T[1:-1,1:-1] = Tn[1:-1,1:-1] + alpha *\
            (dt/dy**2 * (Tn[2:,1:-1] - 2*Tn[1:-1,1:-1] + Tn[:-2,1:-1]) +\
             dt/dx**2 * (Tn[1:-1,2:] - 2*Tn[1:-1,1:-1] + Tn[1:-1,:-2]))

        # Enforce Neumann BCs
        T[-1,:] = T[-2,:]
        T[:,-1] = T[:,-2]

        # Check if we reached T=70C
        if T[j_mid, i_mid] >= 70:
            print ("Center of plate reached 70C at time {0:.2f}s.".format(dt*n))
            break

    if T[j_mid, i_mid]<70:
        print ("Center has not reached 70C yet, it is only {0:.2f}C.".format(T[j_mid, i_mid]))

    return T
```

See the [`break`](https://docs.python.org/2/tutorial/controlflow.html) statement? It exits the `for` loop at the closest time iteration when the plate reaches $70{\rm C}$.

In the code cell below, we define our initial conditions according to the problem set up, and choose the discretization parameters. We start with only 20 spatial steps in each coordinate direction and advance for 500 time steps. You should later experiments with these parameters at your leisure!

```python
L = 1.0e-2
H = 1.0e-2

nx = 21
ny = 21
nt = 500

dx = L/(nx-1)
dy = H/(ny-1)

x = numpy.linspace(0,L,nx)
y = numpy.linspace(0,H,ny)

alpha = 1e-4

Ti = numpy.ones((ny, nx))*20
Ti[0,:]= 100
Ti[:,0] = 100
```

We don't want our solution blowing up, so let's find a time step with $\frac{\alpha \Delta t}{\Delta x^2} = \frac{\alpha \Delta t}{\Delta y^2} = \frac{1}{4}$.

```python
sigma = 0.25
dt = sigma * min(dx, dy)**2 / alpha
T = ftcs(Ti.copy(), nt, alpha, dt, dx, dy)
```

    Center of plate reached 70C at time 0.16s.

### Visualize the results

By now, you're no doubt *very* familiar with the `pyplot.plot` command.  It's great for line plots, scatter plots, etc., but what about when we have two spatial dimensions and another value (temperature) to display?

Are you thinking contour plot?  We're thinking contour plot.  Check out the documentation on [`pyplot.contourf`](http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.contour) (the 'f' denotes "filled" contours).

First, though, we need to create a `meshgrid` purely for plotting purposes:   [`numpy.meshgrid`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.meshgrid.html) is a function that takes two vectors (`x` and `y`, say) and returns two 2D arrays of $x$ and $y$ coordinates that we then use to create the contour plot.

```python
mx, my = numpy.meshgrid(x,y)
```

```python
pyplot.figure(figsize=(8,5))
pyplot.contourf(mx,my,T,20)
pyplot.xlabel('$x$')
pyplot.ylabel('$y$')
pyplot.colorbar();
```

That looks pretty cool!  Note that in the call to `pyplot.contourf` you can specify the number of contour levels to display (we chose 20).  Look at that visualization: does it make physical sense to you, considering that the upper and right sides of the chip are insulated, in our problem?

##### Dig deeper

In the problem we just demonstrated, the chip reaches a temperature of $70{\rm C}$ at a given time, but will it keep increasing? That spells trouble.

Imagine that you have a heat sink instead of an insulator acting on the upper and right sides. What should be the heat flux that the heat sink achieves there, so that the temperature does not exceed $70{\rm C}$ at the center of the chip?

---
###### The cell below loads the style of the notebook

```python
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.css'
HTML(open(css_file, "r").read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Nixie+One' rel='stylesheet' type='text/css'>
<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 750px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
    margin-top:0.8em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Alegreya Sans' sans-serif;
    line-height: 140%;
    font-size: 125%;
    font-weight: 400;
    width:600px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Nixie One', serif;
    font-style:regular;
    font-weight: 400;
    font-size: 45pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h2 {
    font-family: 'Nixie One', serif;
    font-weight: 400;
    font-size: 30pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.1em;
    margin-top: 0.3em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Nixie One', serif;
    margin-top:16px;
    font-size: 22pt;
    font-weight: 600;
    margin-bottom: 3px;
    font-style: regular;
    color: rgb(102,102,0);
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Nixie One', serif;
    font-size: 14pt;
    text-align: center;
    margin-top: 0em;
    margin-bottom: 2em;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Nixie One', sans-serif;
    font-weight: 400;
    font-size: 16pt;
    color: rgb(163,0,0);
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.8em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'PT Mono', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

.CodeMirror{
    font-family: "PT Mono";
    font-size: 90%;
}

</style>