# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # The Basic OLG model

# The OLG model is one of two main analytical frameworks for analyzing the basic intertemporal choice, consumption vs saving. 
#
# OLG is set to capture potential interaction of different generations of individuals in the marketplace hence, providing a tractable alternative to the infinite-horizon economic models

# Imports and set magics:

# +
import numpy as np
from scipy import optimize
import sympy as sm
from sympy import lambdify
import matplotlib.pyplot as plt
import math as m 
from scipy import stats as st
from scipy.optimize import fsolve
import time      
#As we are using a lot of theory in this project, we want the equations to be presented well, why we import the init_session
from sympy import init_session
init_session() 

# autoreload modules when code is run
# %load_ext autoreload
# %autoreload 2

# local modules
import modelproject

# + [markdown] toc-hr-collapsed=true
# # Model description
# -

# In this model, we consider a static economy with an infinite number og households, $i \in N$ 
# In addition, population grows at a constant rate, 
#
# $L_{t}=L_{t-1}(1+n)$
#
#
# Time, t, is discrete and runs to infinity and individuals born at time t live for the period $t$ til $t+1$, meaning the individuals are grouped in two generations, young and old. 
#
# We assume a general, separable utility function for the individuals given by 
#
# $U_t=u(c_{1t})+u(c_{2t+1})(1/(1+\beta))$
#
# $\beta$ is added as a discount factor of the time-horizon. 
#
# Production side is given by competitive firms and a CRS aggregate production function meaning we have the following production criterias
# $Y_{t}=F(K_{t},L_{t})$
#
# $r_{t}=f'(k_{t})$
#
# $w_{t} = f(k_{t})-k_{t}f'(k_{t})$
#
# $\bullet$ NOTE: $k_{t}=K_{t}/L_{t}$
#
# The savings by an individual of a generation is determined as a solution to
#
# $max U_{t}$ s.t.
#
# $w_{t}=s_{t}+c_{1t}$  
#
# $c_{2t+1}=s_{t}(1+r_{t+1})$
#
# Where $c_{1t}$ is consumption of the individual born at t when young at date t and $c_{2t+1}$ consumption when old at date t+1. In this model, the setup is built such that the old individuals will rent their savings of time t as capital to firms at  t+1. Therefore they will eventually receive a gross rate of return given by $(1+r_{t})$. The second consumption constraint incorporates the fact that the individual will consume all his life-earnings. Hence, the lifetime constraint is given by
#
# $w_{t}=c_{1+t}+c_{2t+1}/(1+r_{t+1})$
#
#

# # Solution of the household problem

#
# The solution is found by setting up the Lagrangian
#
# $L(c_{1t},c_{2t+1},\lambda)= u(c_{1t})+u(c_{2t+1})(1/1+\beta)+\lambda[w_{t}-c_{1t}-(c_{2t+1}/(1+r_{t+1}))]$
#
# FOCs:
#
# $u'(c_{1t})=\lambda$
#
# $(1/(1+\beta))u'(c_{t2+1})= \lambda/(1+r_{t+1})$
#
# combining the above to find the Euler equation:
#
# $(1+r_{t+1})/(1+\beta)u'(c_{2t+1}) = u'(c_{1t})$
#
# The interpretation of Euler is that if the consumer gives up one unit of consumption in the young state, the marginal cost will be given by $u'(c_{1t})$ and the marginal benefit will be $(1+r_{t+1})$ units of consumption in next period, when old.
#

# Next, we want to find the **optimal savings path**:
#
# This is done by substituting budget constrain into the Euler eq. above. Hence, we find
#
# $(1+r_{t+1})/(1+\beta)u'(c_{2t+1}) = u'(c_{1t})$ $\rightarrow$ 
#
# $(1+r_{t+1})/(1+\beta)u'((1+r_{t+1})s_{t})=u'(w_{t} - s_{t})$
#
# where optimal savings are a function of wage and the interest rate. 

# Knowing the decision of individuals, we can now examine the **Law of Motion** by aggregating the economy as follows:
#
# As we are working with two periods, we know that the aggregate stock of capital at $t+1$ must equal the total savings of the economy at the time $t$ as well as the non-depreciated capital stock which is carried from $t$ and the 'dissavings' of same period.
#
# This means, the aggregate capital level at time $t+1$ is given by the following:
#
# $K_{t+1}=S_{t}+(1-\delta)K_{t}-K_{t}$ $\rightarrow$
#
# $K_{t+1}=S_{t}-\delta K_{t}$
#
# Remember, $k_{t} \equiv K_{t}/L_{t}$. Hence,
#
# $k_{t+1}(1+n) = s_{t}-\delta k_{t}$
#
# In short, inserting the findings above into the capital accumulation and for simplicity setting $\delta =0$ we now have the law of motion per worker in the entire economy defined by
#
# $k_{t+1} (1+n)=s(f(k_{t})-f'(k_{t})k_{t}, f'(k_{t+1}))$
#
#
# {note: the intuition above is found at http://web.econ.ku.dk/okocg/MAT-OEK/Mak%C3%98k2/Mak%C3%98k2-2016/Forel%C3%A6sninger/Ch%203-2016-1.pdf as well as in our lecture slides from the Macroeconomics course at Copenhagen University}
#

# Now, we wish to solve the model using **sympy**

# +
#First, defining the setting, here the variables we want to use

#Households definitions
beta = sm.symbols('beta')
c1 = sm.symbols('c_{1t}')
c2 = sm.symbols('c_{2t+1}')
lmda = sm.symbols('lambda')
#Savings definitions
r = sm.symbols('r_{t+1}')
w = sm.symbols('w_t')
#Capital stock definitions
s = sm.symbols('s_t')
kt = sm.symbols('k_t')
k1 = sm.symbols('k_{t+1}')
delta = sm.symbols('delta')
n = sm.symbols('n')



# -

# Now, we want to define the utility function. In this case, we use the log utility

# +
#We start by setting a standard utility, then we specify the form afterwards

u_std = sm.Function('u')

def log_u(c):
    return sm.log(c)


def U(c1, c2, beta, u):
    return u(c1) + 1 / (1 + beta) * u(c2)
#Print
U(c1, c2, beta, u = u_std)
# -

# Now, the budget constraints are set up to calculate the intertemporal consumption constraint

# +
#Defining budget constraints
c_1budget = sm.Eq(w, c1 + s)
c_2budget = sm.Eq(c2, (1+r)*s)

#Printing the budget constraints
c_1budget, c_2budget

# -

# Next, we define the intertemporal consumption constraints

#Intertemporal consumption constraint defined
ib = c_1budget.subs(s, sm.solve(c_2budget, s)[0])
#Print intertemporal budget constraint
ib

# Setting up the Lagrangian, deriving FOCs

# +
#Lagrangian
l = U(c1,c2,beta, u = u_std) - lmda *  (ib.rhs - w)

#Print Lagrangian
l

#Deriving FOCs
dc1    = sm.diff(l, c1) 
dc2    = sm.diff(l, c2) 
dlmda = sm.diff(l, lmda)
#Print FOCs
dc1, dc2, dlmda


# -

# Now we find the **Euler Equation**

# +
#Defining the Euler Equation
def euler_equation(dc1, dc2, u, c1, lmda):
    x = dc2.subs(lmda, sm.solve(dc1, lmda)[0])

    euler_eq = sm.Eq(sm.Derivative(u(c1)),
                     sm.solve(x, sm.Derivative(u(c1)))[0]
                    )
    return euler_eq 

#Calling make_euler_equation to calculate the Euler eq
euler = euler_equation(dc1, dc2, u = u_std, c1 = c1, lmda = lmda)

#Printing result
euler

euler = euler.subs(c1, w - s).subs(c2, c_2budget.rhs)
#Printing result
euler

# -

# Euler implicitly determines the value of $s_{t}$. 

# Now we want to determine the development in capital.  

#First we set the functional form of $u_{t}$ using the $log$ function as this enables for further examination. 
euler = euler.replace(u_std, log_u).doit()
#Print
euler


# Now we define the values and setting for the **firms** in the economy.
# To do this, we explore the aggregate production function, normalizing this by using the per capita function as this enables us to better calculate expressions for **interest rates** and **wages** in equilibrium. 

# +
#Production function
Y = sm.symbols('Y_t')
K = sm.symbols('K_t')
L = sm.symbols('L_t')
#including y_t and alpha
y = sm.symbols('y_t')
alpha = sm.symbols('alpha')
#deriving the production function by using the values above

prod_f = sm.Eq(Y,K**alpha * L**(1-alpha))

#Normalizing
normprod_f = sm.Eq(y, kt**alpha)
#Print
normprod_f

# +
ir = sm.Eq(r, sm.Derivative(normprod_f.rhs, kt)).doit()
w_t = sm.Eq(w, normprod_f.rhs - kt*sm.Derivative(normprod_f.rhs, kt)).doit()

#Print results of interest rate and wage
ir, w_t
# -

# Solving for capital period t+1(capital evolution path)

k_1 = sm.Eq(k1, 1/(1+n)* (s - delta*kt) )
#Print 
k_1

# We now derive the **transition path** as follows

# +
k_sav = sm.solve(euler, s)[0].subs(w, w_t.rhs)
# Substituting 
if r in k_sav.atoms():
    k_sav = k_sav.subs(r, ir.rhs)

    
t_path = k_1.subs(s, k_sav)
#Print
t_path
# -

# To make a visualization of the transition path, we now specify a number of parameters and define a function of the transition path

# +
#Defining the transition path
transition_curve = sm.lambdify((kt, beta, delta, n, alpha,r ), t_path.rhs)
def transition_c(kt, beta, delta,alpha, n, r = 0):
    return transition_curve(kt, beta, delta, alpha, n, r)
#Defining equilibrium
eq = sm.lambdify((beta, delta, alpha, n),sm.solve(t_path.subs(k1, kt), kt)[0])
def eqm(beta, delta,alpha, n,):
    return eq(beta, delta, alpha, n)



# -

# Letting alpha vary to show how this affects the transition path, plotting the graph

#Defining variables
ex = np.linspace(0,3,1000)
_b = 0.05
_d = 0.02
_n = 0.2
#Plotting figure
plt.figure(figsize=(14,7))
#Plot 45 degree line
plt.plot(ex, ex, color = 'blue')
#Defining moving function
for _a in np.linspace(0.01, 0.2,5):
    sol = [transition_c(kt = x, alpha = _a, beta = _b, delta = _d, n = _n) for x in ex]
    ks = eqm(alpha = _a, beta = _b, n = _n, delta = _d)
    plt.plot(ex, sol, color = 'red', alpha = 1)    
    plt.annotate(f'$\\alpha$={round(_a,2)}', xy= (ks + 0.01, ks - 0.01))
#Standard settings 
plt.xlabel('$k_t$', size =15)
plt.ylabel('$k_{t+1}$', size=15)
plt.title('Transition curves\n', size = 20)
plt.xlim(0,0.4)
plt.ylim(0,0.4)


# For a lower alpha, the higher is the steady state. The reason behind this, that a higher alpha the more productive is the capital and therefor a higher demand for capital per worker. This means that economy reach the steady state faster. A plausible value would be 1/3 on capital and 1-1/3 n labour. As we can see from our law of motion equation then a higher n (higher population) reduce the capital accumulation and thereby have a lower steady state.

# To further characterize the equilibrium paths, we construct the transition diagram in the plane $(k_{t},k_{t+1})$ examining the convergence towards **steady state**

# +
#Setting variable values
_r = 0.05
_d = 0.05
_a = 0.5

#Defining k steady state
ksteady = 0.0002
ex = range(8)
ex2 = np.linspace(0,1,500)
#Transition
out = list()
for _ in ex:
    ksteady = transition_c(kt = ksteady, alpha = _a, beta = _b, delta = _d, n = _n)
    out.append(ksteady)


# -

# Now we plot the figure in order to illustrate the convergence towards steady state

#
#

res = [transition_c(kt = x, alpha = _a, beta = _b, delta = _d, n = _n) for x in ex2]
#Plotting figure
plt.figure(figsize=(14,7))
#Plot 45 degree line
plt.plot(ex, ex, color = 'blue')
#Plotting the steady state convergence line
plt.plot(ex2, res, color = 'red', alpha = 1)    
plt.step(out[:-1], out[1:], where = 'post', color = 'green', linestyle = '--', alpha = 0.8)
plt.scatter(out[:-1], out[1:], color = 'red', alpha = 0.8)
#Standard settings
plt.xlabel('$k_t$', size = 15)
plt.ylabel('$k_{t+1}$', size = 15)
plt.xlim(0,0.2)
plt.ylim(0,0.2)
plt.title('Convergence towards steady state\n', size = 20)


# Outside the steady-state $k_{t}$ does not grow at a constant rate. Hence, the economy will in time approach the balanced growth path as seen on the graph.
#
#

# # Extended OLG model

# Now we add a government to our model to solve an example of the model using random choosen values. The government collects taxes on both capital and labour. Taxes are given by $\tau$.
#
# We now consider the household problem 
#
# $max$           $u(c_{2t+1})+\beta u(c_{1t})$
#
# s.t.
#
# $c_{2t+1}+s=(1-\tau)w+t_{2t+1}$
#
# $c_{1t}=R+t_{1t}$
#
# where $R$ is the after-tax gross interest rate which is given by
#
# $R=(1+(1-\tau)r)$
#
# Instead of having to substitute and solve for an equation of one unknown, as we did in the beginning, we now want to use the **Gauss-Seidel** method for solving the problem. 
#
# This mdel is an iterative technique for solving a square system of $n$ linear equations with an unknown vaiable, $x$:
#
# $Ax = b$
#
# Which is defined by the iteration:
#
# $Lx^{k+1} =b-Ux^{k}$
#
# Where the variable $x^{k}$ defines the $k^{th}$ approximation of x as well as $x^{k+1}$ is the $k+1$ iteration of x. The matrix, $A$, is decomposed into the component $L$ and a strictly upper triangular component denoted U, meaning $A = L+U$.
#
# The Gauss-Seidel method will solve the lhs of the expression of 
# $Lx^{k+1} =b-Ux^{k}$ which can be written as:
#
# $x^{k+1} =L^{-1} (b-Ux^{k})$
#
# However, as L is of triangular form, the elements of the variable $x^{k+1}$ can be computed sequentially using the forward substitution given by:
#
# $x_i^{k+1} = 1/a_{ii}(b_{i}- \sum_{j=1}^{i-1} a_{ij}x_{j}^{k+1} - \sum_{j=i+1}^{n} a_{ij}x_{j}^{k} ), i=1,...,n$
#
# Using this method, we start with an initial guess of the kapital in time 1t, denoted $K_{guess}$. 
# Next, we solve for the parameters $q, w and R$
#
# $q = \alpha AK^{\alpha-1}_{1t} L^{1-a}$
#
# $w=(1-\alpha)K^{\alpha}_{1t} L^{1-\alpha}$
#
# $R=1+(1-\tau)(q-\delta)$
#
# We then solve for $s*$ which is the optimal savings of the household, given by
#
# $s* =N_{2t+1} ((\beta R((1-\tau)w+t_{2t+1})-t_{1t})/((1+\beta)R))$
#
# Next, we aggregate over all households of the economy to find $K_{new}$ 
#
# $K_{new}=N_{2t+1} \bullet s*$
#
# Lastly, we calculate errors and update capital as follows
#
# $K_{guess} = \lambda K_{new} +(1-\lambda)K_{guess}$ where $0< \lambda <1$ is an updating parameter.
#
# {note: the intuition behind this model is found at https://www.sciencedirect.com/topics/engineering/gauss-seidel-method}

# +
#Set parameters
A_t     = 1
alpha   = 0.33
beta    = 0.8
delta   = 0.0
L_t     = 1
N_2     = 1.0
N_1     = 1.0
tau_L   = 0.22
tau_K   = 0.16
t_2     = 0.0
t_1     = 0.0
#



#Initial guess of capital
error       = 100
guess_lmda  = 0.4   
K_guess     = 0.4
iter        = 1
#Iterating to find the values
while (iter<300) or (error>0.001):
    # Now we want to solve for q and w by
    q = alpha*A_t*K_guess**(alpha-1)
    w = (1-alpha)*A_t*K_guess**alpha
    R = 1 + (1-tau_K)*(q - delta)
    K_new = N_2* (beta*R*((1-tau_L)*w + t_2) - t_1)/((1+beta)*R)
    # Calculate discrepancy between old and new capital stock
    error = abs(K_guess-K_new)/K_guess
    # Update capital stock
    K_guess    = guess_lmda*K_new + (1-guess_lmda)*K_guess
    iter = iter +1


#Results
Ks = K_new
qs = q
Rs = R
rs = qs - delta
ws = w
Ys = A_t*Ks**alpha*L_t**(1-alpha)

# -

# Now we want to find solutions for the entire economy

# +
#Household optimal consumption per period
ss = Ks/N_2
c2s= (1-tau_L)*ws + t_2 - ss
c1s= Rs*ss + t_1

# Residual consumption of the government
Gs = N_2*tau_L*ws + N_1*tau_K*rs*ss

# Finally, aggregate the consumption of the household
Cs = N_2*c2s + N_1*c1s

# To ensure, the above holds, we run a fast check on the condition of the goods market and resoure constraint
ARC = Ys - delta*Ks - Cs - Gs


print("The results using the Gauss-Seidel method are as follows")

print("K* = {:6.4f}".format(Ks))
print("q* = {:6.4f}".format(qs))
print("r* = {:6.4f}".format(rs))
print("R* = {:6.4f}".format(Rs))
print("w* = {:6.4f}".format(ws))
print("Y* = {:6.4f}".format(Ys))
print("-------------------------")
print("ARC = {:6.4f}".format(ARC))
print("Number of iterations = " +str(iter))
print("-------------------------")
print("Optimal consumption of the household is:")
print("s* = ", ss)
print("c1* = ", c1s)
print("c2* = ", c2s)
print("Whereas the residual consumption of the government is:")
print("G* = ", Gs)
print("And the aggregate consumption of the household is:")
print("C* = ", Cs)

# -

#

StarcoderdataJupyterScriptsDedupFiltered