update markov_asset

Author: xuanguang-li

lectures/markov_asset.md

@@ -75,7 +75,10 @@ Let's start with some imports:
 import matplotlib.pyplot as plt
 import numpy as np
 import quantecon as qe
-from numpy.linalg import eigvals, solve
+import jax
+import jax.numpy as jnp
+from jax.numpy.linalg import eigvals, solve
+from typing import NamedTuple
 ```
 
 ## {index}`Pricing Models <single: Pricing Models>`
@@ -91,7 +94,7 @@ Let $\{d_t\}_{t \geq 0}$ be a stream of dividends
 Let's look at some equations that we expect to hold for prices of assets under ex-dividend contracts
 (we will consider cum-dividend pricing in the exercises).
 
-### Risk-Neutral Pricing
+### Risk-neutral pricing
 
 ```{index} single: Pricing Models; Risk-Neutral
 ```
@@ -116,7 +119,7 @@ Here ${\mathbb E}_t [y]$ denotes the best forecast of $y$, conditioned on inform
 
 More precisely, ${\mathbb E}_t [y]$ is the mathematical expectation of $y$ conditional on information available at time $t$.
 
-### Pricing with Random Discount Factor
+### Pricing with random discount factor
 
 ```{index} single: Pricing Models; Risk Aversion
 ```
@@ -145,7 +148,7 @@ This is because such assets pay well when funds are more urgently wanted.
 
 We give examples of how the stochastic discount factor has been modeled below.
 
-### Asset Pricing and Covariances
+### Asset pricing and covariances
 
 Recall that, from the definition of a conditional covariance ${\rm cov}_t (x_{t+1}, y_{t+1})$, we have
 
@@ -174,7 +177,7 @@ Equation {eq}`lteeqs102` asserts that the covariance of the stochastic discount
 
 We give examples of some models of stochastic discount factors that have been proposed later in this lecture and also in a [later lecture](https://python-advanced.quantecon.org/lucas_model.html).
 
-### The Price-Dividend Ratio
+### The price-dividend ratio
 
 Aside from prices, another quantity of interest is the **price-dividend ratio** $v_t := p_t / d_t$.
 
@@ -190,7 +193,7 @@ v_t = {\mathbb E}_t \left[ m_{t+1} \frac{d_{t+1}}{d_t} (1 + v_{t+1}) \right]
 
 Below we'll discuss the implication of this equation.
 
-## Prices in the Risk-Neutral Case
+## Prices in the risk-neutral case
 
 What can we say about price dynamics on the basis of the models described above?
 
@@ -203,7 +206,7 @@ For now we'll study  the risk-neutral case in which  the stochastic discount fac
 
 We'll  focus on how an asset  price depends on a dividend process.
 
-### Example 1: Constant Dividends
+### Example 1: constant dividends
 
 The simplest case is risk-neutral price of a constant, non-random dividend stream $d_t = d > 0$.
 
@@ -234,7 +237,7 @@ This is the equilibrium price in the constant dividend case.
 Indeed, simple algebra shows that setting $p_t = \bar p$ for all $t$
 satisfies the difference equation $p_t = \beta (d + p_{t+1})$.
 
-### Example 2: Dividends with Deterministic Growth Paths
+### Example 2: dividends with deterministic growth paths
 
 Consider a growing, non-random dividend process $d_{t+1} = g d_t$
 where $0 < g \beta < 1$.
@@ -267,7 +270,7 @@ $$
 This is called the *Gordon formula*.
 
 (mass_mg)=
-### Example 3: Markov Growth, Risk-Neutral Pricing
+### Example 3: Markov growth, risk-neutral pricing
 
 Next, we consider a dividend process
 
@@ -310,14 +313,14 @@ The next figure shows a simulation, where
 * $\{X_t\}$ evolves as a discretized AR1 process produced using {ref}`Tauchen's method <fm_ex3>`.
 * $g_t = \exp(X_t)$, so that $\ln g_t = X_t$ is the growth rate.
 
-```{code-cell} ipython
+```{code-cell} ipython3
 n = 7
 mc = qe.tauchen(n, 0.96, 0.25)
 sim_length = 80
 
-x_series = mc.simulate(sim_length, init=np.median(mc.state_values))
-g_series = np.exp(x_series)
-d_series = np.cumprod(g_series) # Assumes d_0 = 1
+x_series = mc.simulate(sim_length, init=jnp.median(mc.state_values))
+g_series = jnp.exp(x_series)
+d_series = jnp.cumprod(g_series) # Assumes d_0 = 1
 
 series = [x_series, g_series, d_series, np.log(d_series)]
 labels = ['$X_t$', '$g_t$', '$d_t$', r'$\log \, d_t$']
@@ -330,7 +333,7 @@ plt.tight_layout()
 plt.show()
 ```
 
-#### Pricing Formula
+#### Pricing formula
 
 To obtain asset prices in this setting, let's adapt our analysis from the case of deterministic growth.
 
@@ -400,18 +403,18 @@ As before, we'll generate $\{X_t\}$  as a {ref}`discretized AR1 process <fm_ex3>
 
 Here's the code, including a test of the spectral radius condition
 
-```{code-cell} python3
+```{code-cell} ipython3
 n = 25  # Size of state space
 β = 0.9
 mc = qe.tauchen(n, 0.96, 0.02)
 
-K = mc.P * np.exp(mc.state_values)
+K = mc.P * jnp.exp(mc.state_values)
 
 warning_message = "Spectral radius condition fails"
-assert np.max(np.abs(eigvals(K))) < 1 / β,  warning_message
+assert jnp.max(jnp.abs(eigvals(K))) < 1 / β,  warning_message
 
-I = np.identity(n)
-v = solve(I - β * K, β * K @ np.ones(n))
+I = jnp.identity(n)
+v = solve(I - β * K, β * K @ jnp.ones(n))
 
 fig, ax = plt.subplots()
 ax.plot(mc.state_values, v, 'g-o', lw=2, alpha=0.7, label='$v$')
@@ -440,7 +443,7 @@ We'll price several distinct assets, including
 * A consol (a type of bond issued by the UK government in the 19th century)
 * Call options on a consol
 
-### Pricing a Lucas Tree
+### Pricing a Lucas tree
 
 ```{index} single: Finite Markov Asset Pricing; Lucas Tree
 ```
@@ -539,46 +542,51 @@ v = (I - \beta J)^{-1} \beta  J {\mathbb 1}
 We will define a function tree_price to compute $v$ given parameters stored in
 the class AssetPriceModel
 
-```{code-cell} python3
-class AssetPriceModel:
+```{code-cell} ipython3
+class AssetPriceModel(NamedTuple):
     """
     A class that stores the primitives of the asset pricing model.
 
     Parameters
     ----------
-    β : scalar, float
-        Discount factor
     mc : MarkovChain
-        Contains the transition matrix and set of state values for the state
-        process
-    γ : scalar(float)
-        Coefficient of risk aversion
+        Contains the transition matrix and set of state values
     g : callable
         The function mapping states to growth rates
-
+    β : float
+        Discount factor
+    γ : float
+        Coefficient of risk aversion
+    n: int
+        The number of states
+    """
+    mc: qe.MarkovChain
+    g: callable
+    β: float
+    γ: float
+    n: int
+    
+
+def create_ap_model(mc=None, g=jnp.exp, β=0.96, γ=2.0):
+    """Create an AssetPriceModel class"""
+    if mc is None:
+        n, ρ, σ = 25, 0.9, 0.02
+        mc = qe.tauchen(n, ρ, σ)
+    else:
+        mc = mc
+        n = mc.P.shape[0]
+
+    return AssetPriceModel(mc=mc, g=g, β=β, γ=γ, n=n)
+
+
+def test_stability(Q, β):
     """
-    def __init__(self, β=0.96, mc=None, γ=2.0, g=np.exp):
-        self.β, self.γ = β, γ
-        self.g = g
-
-        # A default process for the Markov chain
-        if mc is None:
-            self.ρ = 0.9
-            self.σ = 0.02
-            self.mc = qe.tauchen(n, self.ρ, self.σ)
-        else:
-            self.mc = mc
-
-        self.n = self.mc.P.shape[0]
-
-    def test_stability(self, Q):
-        """
-        Stability test for a given matrix Q.
-        """
-        sr = np.max(np.abs(eigvals(Q)))
-        if not sr < 1 / self.β:
-            msg = f"Spectral radius condition failed with radius = {sr}"
-            raise ValueError(msg)
+    Stability test for a given matrix Q.
+    """
+    sr = np.max(np.abs(eigvals(Q)))
+    if not sr < 1 / β:
+        msg = f"Spectral radius condition failed with radius = {sr}"
+        raise ValueError(msg)
 
 
 def tree_price(ap):
@@ -601,11 +609,11 @@ def tree_price(ap):
     J = P * ap.g(y)**(1 - γ)
 
     # Make sure that a unique solution exists
-    ap.test_stability(J)
+    test_stability(J, β)
 
     # Compute v
-    I = np.identity(ap.n)
-    Ones = np.ones(ap.n)
+    I = jnp.identity(ap.n)
+    Ones = jnp.ones(ap.n)
     v = solve(I - β * J, β * J @ Ones)
 
     return v
@@ -614,16 +622,16 @@ def tree_price(ap):
 Here's a plot of $v$ as a function of the state for several values of $\gamma$,
 with a positively correlated Markov process and $g(x) = \exp(x)$
 
-```{code-cell} python3
+```{code-cell} ipython3
 γs = [1.2, 1.4, 1.6, 1.8, 2.0]
-ap = AssetPriceModel()
+ap = create_ap_model()
 states = ap.mc.state_values
 
 fig, ax = plt.subplots()
 
 for γ in γs:
-    ap.γ = γ
-    v = tree_price(ap)
+    tem_ap = create_ap_model(mc=ap.mc, g=ap.g, β=ap.β, γ=γ)
+    v = tree_price(tem_ap)
     ax.plot(states, v, lw=2, alpha=0.6, label=rf"$\gamma = {γ}$")
 
 ax.set_title('Price-dividend ratio as a function of the state')
@@ -706,7 +714,7 @@ p = (I - \beta M)^{-1} \beta M \zeta {\mathbb 1}
 
 The above is implemented in the function consol_price.
 
-```{code-cell} python3
+```{code-cell} ipython3
 def consol_price(ap, ζ):
     """
     Computes price of a consol bond with payoff ζ
@@ -715,26 +723,25 @@ def consol_price(ap, ζ):
     ----------
     ap: AssetPriceModel
         An instance of AssetPriceModel containing primitives
-
+    
     ζ : scalar(float)
         Coupon of the console
 
     Returns
     -------
     p : array_like(float)
         Console bond prices
-
     """
     # Simplify names, set up matrices
     β, γ, P, y = ap.β, ap.γ, ap.mc.P, ap.mc.state_values
     M = P * ap.g(y)**(- γ)
 
     # Make sure that a unique solution exists
-    ap.test_stability(M)
+    test_stability(M, β)
 
     # Compute price
-    I = np.identity(ap.n)
-    Ones = np.ones(ap.n)
+    I = jnp.identity(ap.n)
+    Ones = jnp.ones(ap.n)
     p = solve(I - β * M, β * ζ * M @ Ones)
 
     return p
@@ -812,7 +819,7 @@ Start at some initial $w$ and iterate with $T$ to convergence .
 
 We can find the solution with the following function call_option
 
-```{code-cell} python3
+```{code-cell} ipython3
 def call_option(ap, ζ, p_s, ϵ=1e-7):
     """
     Computes price of a call option on a consol bond.
@@ -828,7 +835,7 @@ def call_option(ap, ζ, p_s, ϵ=1e-7):
     p_s : scalar(float)
         Strike price
 
-    ϵ : scalar(float), optional(default=1e-8)
+    ϵ : scalar(float), optional(default=1e-7)
         Tolerance for infinite horizon problem
 
     Returns
@@ -842,26 +849,35 @@ def call_option(ap, ζ, p_s, ϵ=1e-7):
     M = P * ap.g(y)**(- γ)
 
     # Make sure that a unique consol price exists
-    ap.test_stability(M)
+    test_stability(M, β)
 
     # Compute option price
     p = consol_price(ap, ζ)
-    w = np.zeros(ap.n)
+    w = jnp.zeros(ap.n)
     error = ϵ + 1
-    while error > ϵ:
+
+    def step(state):
+        w, error = state
         # Maximize across columns
-        w_new = np.maximum(β * M @ w, p - p_s)
+        w_new = jnp.maximum(β * M @ w, p - p_s)
         # Find maximal difference of each component and update
-        error = np.amax(np.abs(w - w_new))
-        w = w_new
+        error_new = jnp.amax(jnp.abs(w - w_new))
+        return (w_new, error_new)
 
-    return w
+    # Check whether converged
+    def cond(state):
+        _, error = state
+        return error > ϵ
+
+    final_w, _ = jax.lax.while_loop(cond, step, (w, error))
+
+    return final_w
 ```
 
 Here's a plot of $w$ compared to the consol price when $P_S = 40$
 
-```{code-cell} python3
-ap = AssetPriceModel(β=0.9)
+```{code-cell} ipython3
+ap = create_ap_model(β=0.9)
 ζ = 1.0
 strike_price = 40