Fix typos and np/jnp issues (#135)

Author: kp992

lectures/markov_asset.md

@@ -4,7 +4,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.5
+    jupytext_version: 1.16.1
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -31,12 +31,10 @@ you can jump to [](eq:ntecx2).
 
 The code outputs below are generated by machine connected to the following GPU
 
-
 ```{code-cell} ipython3
 !nvidia-smi
 ```
 
-
 In addition to what's in Anaconda, this lecture will need the following libraries:
 
 ```{code-cell} ipython3
@@ -410,7 +408,7 @@ def create_model(N=100,         # size of state space for Markov chain
                  σ=0.01,        # persistence parameter for Markov chain
                  β=0.98,        # discount factor
                  γ=2.5,         # coefficient of risk aversion
-                 μ_c=0.01,      # mean growth of consumtion
+                 μ_c=0.01,      # mean growth of consumption
                  μ_d=0.01,      # mean growth of dividends
                  σ_c=0.02,      # consumption volatility
                  σ_d=0.04):     # dividend volatility
@@ -431,14 +429,14 @@ Here's a function that does this using loops.
 
 ```{code-cell} ipython3
 def compute_K_loop(model):
-    # Setp up
+    # unpack
     P, S, β, γ, μ_c, μ_d, σ_c, σ_d = model
     N = len(S)
     K = np.empty((N, N))
     a = μ_d - γ * μ_c
     for i, x in enumerate(S):
         for j, y in enumerate(S):
-            e = jnp.exp(a + (1 - γ) * x + (σ_d**2 + γ**2 * σ_c**2) / 2)
+            e = np.exp(a + (1 - γ) * x + (σ_d**2 + γ**2 * σ_c**2) / 2)
             K[i, j] = β * e * P[i, j]
     return K
 ```
@@ -447,24 +445,24 @@ To exploit the parallelization capabilities of JAX, let's also write a vectorize
 
 ```{code-cell} ipython3
 def compute_K(model):
-    # Setp up
+    # unpack
     P, S, β, γ, μ_c, μ_d, σ_c, σ_d = model
     N = len(S)
     # Reshape and multiply pointwise using broadcasting
-    x = jnp.reshape(S, (N, 1))
+    x = np.reshape(S, (N, 1))
     a = μ_d - γ * μ_c
-    e = jnp.exp(a + (1 - γ) * x + (σ_d**2 + γ**2 * σ_c**2) / 2)
+    e = np.exp(a + (1 - γ) * x + (σ_d**2 + γ**2 * σ_c**2) / 2)
     K = β * e * P
     return K
 ```
 
-These two functions produce the saem output:
+These two functions produce the same output:
 
 ```{code-cell} ipython3
 model = create_model(N=10)
 K1 = compute_K(model)
 K2 = compute_K_loop(model)
-jnp.allclose(K1, K2)
+np.allclose(K1, K2)
 ```
 
 Now we can compute the price-dividend ratio:
@@ -492,9 +490,9 @@ def price_dividend_ratio(model, test_stable=True):
         test_stability(K)
 
     # Compute v
-    I = jnp.identity(N)
-    ones_vec = jnp.ones(N)
-    v = jnp.linalg.solve(I - K, K @ ones_vec)
+    I = np.identity(N)
+    ones_vec = np.ones(N)
+    v = np.linalg.solve(I - K, K @ ones_vec)
 
     return v
 ```
@@ -504,7 +502,7 @@ Here's a plot of $v$ as a function of the state for several values of $\gamma$.
 ```{code-cell} ipython3
 model = create_model()
 S = model.S
-γs = jnp.linspace(2.0, 3.0, 5)
+γs = np.linspace(2.0, 3.0, 5)
 
 fig, ax = plt.subplots()
 
@@ -701,7 +699,7 @@ def create_sv_model(β=0.98,        # discount factor
                     bar_σ=0.01,    # volatility scaling parameter
                     ρ_z=0.9,       # persistence parameter for z
                     σ_z=0.01,      # persistence parameter for z
-                    μ_c=0.001,     # mean growth of consumtion
+                    μ_c=0.001,     # mean growth of consumption
                     μ_d=0.005):    # mean growth of dividends
 
     mc = qe.tauchen(I, ρ_c, σ_c)
@@ -763,7 +761,7 @@ def sv_pd_ratio(sv_model, test_stable=True):
         price-dividend ratio
 
     """
-    # Setp up
+    # unpack
     P, hc_grid, Q, hd_grid, R, z_grid, β, γ, bar_σ, μ_c, μ_d = sv_model
     I, J, K = len(hc_grid), len(hd_grid), len(z_grid)
     N = I * J * K
@@ -862,12 +860,12 @@ def compute_A_jax(sv_model, shapes):
     I, J, K = shapes
     N = I * J * K
     # Reshape and broadcast over (i, j, k, i', j', k')
-    hc = np.reshape(hc_grid,     (I, 1, 1, 1,  1,  1))
-    hd = np.reshape(hd_grid,     (1, J, 1, 1,  1,  1))
-    z = np.reshape(z_grid,       (1, 1, K, 1,  1,  1))
-    P = np.reshape(P,            (I, 1, 1, I,  1,  1))
-    Q = np.reshape(Q,            (1, J, 1, 1,  J,  1))
-    R = np.reshape(R,            (1, 1, K, 1,  1,  K))
+    hc = jnp.reshape(hc_grid,     (I, 1, 1, 1,  1,  1))
+    hd = jnp.reshape(hd_grid,     (1, J, 1, 1,  1,  1))
+    z = jnp.reshape(z_grid,       (1, 1, K, 1,  1,  1))
+    P = jnp.reshape(P,            (I, 1, 1, I,  1,  1))
+    Q = jnp.reshape(Q,            (1, J, 1, 1,  J,  1))
+    R = jnp.reshape(R,            (1, 1, K, 1,  1,  K))
     # Compute A and then reshape to create a matrix
     a = μ_d - γ * μ_c
     b = bar_σ**2 * (jnp.exp(2 * hd) + γ**2 * jnp.exp(2 * hc)) / 2
@@ -896,7 +894,7 @@ def sv_pd_ratio_jax(sv_model, shapes):
         price-dividend ratio
 
     """
-    # Setp up
+    # unpack
     P, hc_grid, Q, hd_grid, R, z_grid, β, γ, bar_σ, μ_c, μ_d = sv_model
     I, J, K = len(hc_grid), len(hd_grid), len(z_grid)
     shapes = I, J, K
@@ -971,13 +969,13 @@ def A(g, sv_model, shapes):
     P, hc_grid, Q, hd_grid, R, z_grid, β, γ, bar_σ, μ_c, μ_d = sv_model
     I, J, K = shapes
     # Reshape and broadcast over (i, j, k, i', j', k')
-    hc = np.reshape(hc_grid,     (I, 1, 1, 1,  1,  1))
-    hd = np.reshape(hd_grid,     (1, J, 1, 1,  1,  1))
-    z = np.reshape(z_grid,       (1, 1, K, 1,  1,  1))
-    P = np.reshape(P,            (I, 1, 1, I,  1,  1))
-    Q = np.reshape(Q,            (1, J, 1, 1,  J,  1))
-    R = np.reshape(R,            (1, 1, K, 1,  1,  K))
-    g = np.reshape(g,            (1, 1, 1, I,  J,  K))
+    hc = jnp.reshape(hc_grid,     (I, 1, 1, 1,  1,  1))
+    hd = jnp.reshape(hd_grid,     (1, J, 1, 1,  1,  1))
+    z = jnp.reshape(z_grid,       (1, 1, K, 1,  1,  1))
+    P = jnp.reshape(P,            (I, 1, 1, I,  1,  1))
+    Q = jnp.reshape(Q,            (1, J, 1, 1,  J,  1))
+    R = jnp.reshape(R,            (1, 1, K, 1,  1,  K))
+    g = jnp.reshape(g,            (1, 1, 1, I,  J,  K))
     a = μ_d - γ * μ_c
     b = bar_σ**2 * (jnp.exp(2 * hd) + γ**2 * jnp.exp(2 * hc)) / 2
     κ = jnp.exp(a + (1 - γ) * z + b)
@@ -991,12 +989,10 @@ that acts directly on the linear operator `A`.
 
 ```{code-cell} ipython3
 def sv_pd_ratio_linop(sv_model, shapes):
-
-    # Setp up
     P, hc_grid, Q, hd_grid, R, z_grid, β, γ, bar_σ, μ_c, μ_d = sv_model
     I, J, K = shapes
 
-    ones_array = np.ones((I, J, K))
+    ones_array = jnp.ones((I, J, K))
     # Set up the operator g -> (I - A) g
     J = lambda g: g - A(g, sv_model, shapes)
     # Solve v = (I - A)^{-1} A 1