update code

Author: HumphreyYang

sandpit/cons_smooth_tom_v2.md

@@ -4,7 +4,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.4
+    jupytext_version: 1.14.5
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -22,6 +22,7 @@ In this notebook, we'll present  some useful models of economic dynamics using o
 ```{code-cell} ipython3
 import numpy as np
 import matplotlib.pyplot as plt
+from collections import namedtuple
 ```
 
 +++ {"user_expressions": []}
@@ -70,14 +71,29 @@ Below, we'll describe how to execute these steps using linear algebra -- matrix
 
 We shall eventually evaluate alternative budget feasible consumption paths $\vec c$ using the following **welfare criterion**
 
-$$
+```{math}
+:label: welfare
+
 W = \sum_{t=0}^T \beta^t (g_1 c_t - \frac{g_2}{2} c_t^2 )
-$$
+```
 
 where $g_1 > 0, g_2 > 0$.  
 
 We shall see that when $\beta R = 1$ (a condition assumed by Milton Friedman and Robert Hall), this criterion assigns higher welfare to **smoother** consumption paths.
 
+Here we use default parameters $R = 1.05$, $g_1 = 1$, $g_2 = 1/2$, and $T = 65$. 
+
+We create a namedtuple to store these parameters with default values.
+
+```{code-cell} ipython3
+ConsumptionSmoothing = namedtuple("ConsumptionSmoothing", ["R", "g1", "g2", "β_seq", "T"])
+
+def creat_cs_model(R=1.05, g1=1, g2=1/2, T=65):
+    β = 1/R
+    β_seq = np.array([β**i for i in range(T+1)])
+    return ConsumptionSmoothing(R=1.05, g1=1, g2=1/2, β_seq=β_seq, T=65)
+```
+
 +++ {"user_expressions": []}
 
 ## Difference equations with linear algebra ##
@@ -209,7 +225,29 @@ $$
 
 This is the consumption-smoothing model in a nutshell.
 
-We'll put the model through some paces with Python code below.
+We implement this model in `compute_optimal`
+
+```{code-cell} ipython3
+def compute_optimal(model, a0, y_seq):
+    R, T = model.R, model.T
+
+    # non-financial wealth
+    h0 = model.β_seq @ y_seq     # since β = 1/R
+
+    # c0
+    c0 = (1 - 1/R) / (1 - (1/R)**(T+1)) * (a0 + h0)
+    c_seq = c0*np.ones(T+1)
+
+    # verify
+    A = np.diag(-R*np.ones(T), k=-1) + np.eye(T+1)
+    b = y_seq - c_seq
+    b[0] = b[0] + a0
+
+    a_seq = np.linalg.inv(A) @ b
+    a_seq = np.concatenate([[a0], a_seq])
+
+    return c_seq, a_seq
+```
 
 +++ {"user_expressions": []}
 
@@ -218,13 +256,11 @@ We'll put the model through some paces with Python code below.
 As promised, we'll provide step by step instructions on how to use linear algebra, readily implemented
 in Python, to solve the consumption smoothing model.
 
-**Note to programmer teammate:**
-
 In the calculations below, please we'll  set default values of  $R > 1$, e.g., $R = 1.05$, and $\beta = R^{-1}$.
 
 #### Step 1 ####
 
-For some $T+1 \times 1$ $y$ vector, use matrix algebra to compute 
+For some $(T+1) \times 1$ $y$ vector, use matrix algebra to compute 
 
 $$
 \sum_{t=0}^T R^{-t} y_t = \begin{bmatrix} 1 & R^{-1} & \cdots & R^{-T} \end{bmatrix}
@@ -275,8 +311,53 @@ $$
 
 Let's verify this with our Python code.
 
+We use an example where the consumer inherits $a_0<0$ (which can be interpreted as a student debt).
+
+The income process $\{y_t\}_{t=0}^{T}$ is constant and positive up to $t=45$ and then becomes zero afterward.
+
+```{code-cell} ipython3
+# Financial wealth
+a0 = -2     # such as "student debt"
+
+# Income process
+y_seq = np.concatenate([np.ones(46), np.zeros(20)])
+
+cs_model = creat_cs_model()
+c_seq, a_seq = compute_optimal(cs_model, a0, y_seq)
+
+print('check a_T+1=0:', np.abs(a_seq[-1] - 0) <= 1e-8)
+```
+
+```{code-cell} ipython3
+# Sequence Length
+T = cs_model.T
+
+plt.plot(range(T+1), y_seq, label='income')
+plt.plot(range(T+1), c_seq, label='consumption')
+plt.plot(range(T+2), a_seq, label='asset')
+plt.plot(range(T+2), np.zeros(T+2), '--')
+
+plt.legend()
+plt.xlabel(r'$t$')
+plt.ylabel(r'$c_t,y_t,a_t$')
+plt.show()
+```
+
++++ {"user_expressions": []}
+
+We can evaluate the welfare using the formula {numref}`welfare`
+
+```{code-cell} ipython3
+def welfare(model, c_seq):
+    β_seq, g1, g2 = model.β_seq, model.g1, model.g2
+
+    u_seq = g1 * c_seq - g2/2 * c_seq**2
+    return β_seq @ u_seq
 
+print('Welfare:', welfare(cs_model, c_seq))
+```
 
++++ {"user_expressions": []}
 
 ### Feasible consumption variations ###
 
@@ -331,158 +412,79 @@ Given $R$, we thus have a two parameter class of budget feasible variations $\ve
 to compute alternative consumption paths, then evaluate their welfare.
 
 **Note to John:** We can do some fun simple experiments with these variations -- we can use
-graphs to show that, when $\beta R =1$ and  starting from the smooth path, all nontrivial budget-feasible variations lower welfare according to the criterion above.  
-
-We can even use the Python numpy grad command to compute derivatives of welfare with respect to our two parameters.  
-
-We are teaching the key idea beneath the **calculus of variations**.
-
-```{code-cell} ipython3
-class Consumption_smoothing:
-    "A class of the Permanent Income model of consumption"
-    
-    def __init__(self, R, y_seq, a0, g1, g2, T):
-        self.a0, self.y_seq, self.R, self.β = a0, y_seq, R, 1/R    # set β = 1/R
-        self.g1, self.g2 = g1, g2       # welfare parameter
-        self.T = T
-        
-        self.β_seq = np.array([self.β**i for i in range(T+1)])
-        
-    def compute_optimal(self, verbose=1):
-        R, y_seq, a0, T = self.R, self.y_seq, self.a0, self.T
-        
-        # non-financial wealth
-        h0 = self.β_seq @ y_seq     # since β = 1/R
-        
-        # c0
-        c0 = (1 - 1/R) / (1 - (1/R)**(T+1)) * (a0 + h0)
-        c_seq = c0*np.ones(T+1)
-        
-        # verify
-        A = np.diag(-R*np.ones(T), k=-1) + np.eye(T+1)
-        b = y_seq - c_seq
-        b[0] = b[0] + a0
-        
-        a_seq = np.linalg.inv(A) @ b
-        a_seq = np.concatenate([[a0], a_seq])
-        
-        # check that a_T+1 = 0
-        if verbose==1:
-            print('check a_T+1=0:', np.abs(a_seq[-1] - 0) <= 1e-8)
-        
-        return c_seq, a_seq
-    
-    def welfare(self, c_seq):
-        β_seq, g1, g2 = self.β_seq, self.g1, self.g2
-        
-        u_seq = g1 * c_seq - g2/2 * c_seq**2
-        return β_seq @ u_seq
-        
-    
-    def compute_variation(self, ξ1, ϕ, verbose=1):
-        R, T, β_seq = self.R, self.T, self.β_seq
-        
-        ξ0 = ξ1*((1 - 1/R) / (1 - (1/R)**(T+1))) * ((1 - (ϕ/R)**(T+1)) / (1 - ϕ/R))
-        v_seq = np.array([(ξ1*ϕ**t - ξ0) for t in range(T+1)])
-        
-        # check if it is feasible
-        if verbose==1:
-            print('check feasible:', np.round(β_seq @ v_seq, 7)==0)     # since β = 1/R
-        
-        c_opt, _ = self.compute_optimal(verbose=verbose)
-        cvar_seq = c_opt + v_seq
-        
-        return cvar_seq
-```
+graphs to show that, when $\beta R = 1$ and  starting from the smooth path, all nontrivial budget-feasible variations lower welfare according to the criterion above.  
 
-+++ {"user_expressions": []}
-
-Below is an example where the consumer inherits $a_0<0$ (which can be interpreted as a student debt).
-
-The income process $\{y_t\}_{t=0}^{T}$ is constant and positive up to $t=45$ and then becomes zero afterward.
+Now let's compute and visualize the variations
 
 ```{code-cell} ipython3
-# parameters
-T=65
-R = 1.05
-g1 = 1
-g2 = 1/2
-
-# financial wealth
-a0 = -2     # such as "student debt"
+def compute_variation(model, ξ1, ϕ, a0, y_seq, verbose=1):
+    R, T, β_seq = model.R, model.T, model.β_seq
 
-# income process
-y_seq = np.concatenate([np.ones(46), np.zeros(20)])
-
-# create an instance
-mc = Consumption_smoothing(R=R, y_seq=y_seq, a0=a0, g1=g1, g2=g2, T=T)
-c_seq, a_seq = mc.compute_optimal()
-
-# compute welfare 
-print('Welfare:', mc.welfare(c_seq))
-```
+    ξ0 = ξ1*((1 - 1/R) / (1 - (1/R)**(T+1))) * ((1 - (ϕ/R)**(T+1)) / (1 - ϕ/R))
+    v_seq = np.array([(ξ1*ϕ**t - ξ0) for t in range(T+1)])
+    
+    if verbose == 1:
+        print('check feasible:', np.isclose(β_seq @ v_seq, 0))     # since β = 1/R
 
-```{code-cell} ipython3
-plt.plot(range(T+1), y_seq, label='income')
-plt.plot(range(T+1), c_seq, label='consumption')
-plt.plot(range(T+2), a_seq, label='asset')
-plt.plot(range(T+2), np.zeros(T+2), '--')
+    c_opt, _ = compute_optimal(model, a0, y_seq)
+    cvar_seq = c_opt + v_seq
 
-plt.legend()
-plt.xlabel(r'$t$')
-plt.ylabel(r'$c_t,y_t,a_t$')
-plt.show()
+    return cvar_seq
 ```
 
 +++ {"user_expressions": []}
 
-We can visualize how $\xi_1$ and $\phi$ controls **budget-feasible variations**.
-
-```{code-cell} ipython3
-# visualize variational paths
-cvar_seq1 = mc.compute_variation(ξ1=.01, ϕ=.95)
-cvar_seq2 = mc.compute_variation(ξ1=.05, ϕ=.95)
-cvar_seq3 = mc.compute_variation(ξ1=.01, ϕ=1.02)
-cvar_seq4 = mc.compute_variation(ξ1=.05, ϕ=1.02)
-```
+We visualize variations with $\xi_1 \in \{.01, .05\}$ and $\phi \in \{.95, 1.02\}$
 
 ```{code-cell} ipython3
-print('welfare of optimal c: ', mc.welfare(c_seq))
-print('variation 1: ', mc.welfare(cvar_seq1))
-print('variation 2:', mc.welfare(cvar_seq2))
-print('variation 3: ', mc.welfare(cvar_seq3))
-print('variation 4:', mc.welfare(cvar_seq4))
-```
+fig, ax = plt.subplots()
+
+ξ1s = [.01, .05]
+ϕs= [.95, 1.02]
+colors = {.01: 'tab:blue', .05: 'tab:green'}
+
+params = np.array(np.meshgrid(ξ1s, ϕs)).T.reshape(-1, 2)
+
+for i, param in enumerate(params):
+    ξ1, ϕ = param
+    print(f'variation {i}: ξ1={ξ1}, ϕ={ϕ}')
+    cvar_seq = compute_variation(model=cs_model, 
+                                 ξ1=ξ1, ϕ=ϕ, a0=a0, 
+                                 y_seq=y_seq)
+    print(f'welfare={welfare(cs_model, cvar_seq)}')
+    print('-'*64)
+    if i % 2 == 0:
+        ls = '-.'
+    else: 
+        ls = '-'  
+    ax.plot(range(T+1), cvar_seq, ls=ls, color=colors[ξ1], label=fr'$\xi_1 = {ξ1}, \phi = {ϕ}$')
 
-```{code-cell} ipython3
 plt.plot(range(T+1), c_seq, color='orange', label=r'Optimal $\vec{c}$ ')
-plt.plot(range(T+1), cvar_seq1, color='tab:blue', label=r'$\xi_1 = 0.01, \phi = 0.95$')
-plt.plot(range(T+1), cvar_seq2, color='tab:blue', ls='-.', label=r'$\xi_1 = 0.05, \phi = 0.95$')
-plt.plot(range(T+1), cvar_seq3, color='tab:green', label=r'$\xi_1 = 0.01, \phi = 1.02$')
-plt.plot(range(T+1), cvar_seq4, color='tab:green', ls='-.', label=r'$\xi_1 = 0.05, \phi = 1.02$')
-
 
 plt.legend()
 plt.xlabel(r'$t$')
 plt.ylabel(r'$c_t$')
 plt.show()
 ```
 
++++ {"user_expressions": []}
+
+We can even use the Python numpy grad command to compute derivatives of welfare with respect to our two parameters.  
+
+We are teaching the key idea beneath the **calculus of variations**.
+
 ```{code-cell} ipython3
-def welfare_ϕ(mc, ξ1, ϕ):
-    "Compute welfare of variation sequence for given ϕ, ξ1 with an instance of our model mc"
-    cvar_seq = mc.compute_variation(ξ1=ξ1, ϕ=ϕ, verbose=0)
-    return mc.welfare(cvar_seq)
+def welfare_ϕ(ξ1, ϕ):
+    "Compute welfare of variation sequence for given ϕ, ξ1 with a consumption smoothing model"
+    cvar_seq = compute_variation(cs_model, ξ1=ξ1, ϕ=ϕ, a0=a0, 
+                                 y_seq=y_seq, verbose=0)
+    return welfare(cs_model, cvar_seq)
 
-welfare_φ = np.vectorize(welfare_φ)
+welfare_ϕ_vec = np.vectorize(welfare_ϕ)
 ξ1_arr = np.linspace(-0.5, 0.5, 20)
 
-plt.plot(ξ1_arr, welfare_φ(mc, ξ1=ξ1_arr , ϕ=1.02))
+plt.plot(ξ1_arr, welfare_ϕ_vec(ξ1_arr, 1.02))
 plt.ylabel('welfare')
 plt.xlabel(r'$\xi_1$')
 plt.show()
 ```
-
-```{code-cell} ipython3
-
-```