```{code-cell} ipython3

import matplotlib.pyplot as plt

import numpy as np

from quantecon import MarkovChain

import quantecon as qe

import jax

import jax.numpy as jnp

from jax import vmap

from typing import NamedTuple

from functools import partial

```

When $\{R_t\}$ was constant we required that $\beta R < 1$.

Since it is now stochastic, we require that

Since it is now stochastic, we require (see {cite}`ma2020income`) that

```{math}

:label: fpbc2

\left(\mathbb E \prod_{t=1}^n R_t \right)^{1/n}

```

Notice that, when $\{R_t\}$ takes some constant value $R$, this

reduces to the previous restriction $\beta R < 1$.

The value $G_R$ can be thought of as the long run (geometric) average

gross rate of return.

More intuition behind {eq}`fpbc2` is provided in {cite}`ma2020income`.

To simplify this lecture, we will *assume that the interest rate process is

IID*.

In that case, it is clear from the definition of $G_R$ that $G_R$ is just $\mathbb E R_t$.

Discussion on how to check it is given below.

We test the condition $\beta \mathbb E R_t < 1$ in the code below.

Finally, we impose some routine technical restrictions on non-financial income.

In what follows, we use linear interpolation.

### Testing the Assumptions

Convergence of time iteration is dependent on the condition $\beta G_R < 1$ being satisfied.

One can check this using the fact that $G_R$ is equal to the spectral

radius of the matrix $L$ defined by

$$

L(z, \hat z) := P(z, \hat z) \int R(\hat z, x) \phi(x) dx

$$

This identity is proved in {cite}`ma2020income`, where $\phi$ is the

density of the innovation $\zeta_t$ to returns on assets.

(Remember that $\mathsf Z$ is a finite set, so this expression defines a matrix.)

Checking the condition is even easier when $\{R_t\}$ is IID.

In that case, it is clear from the definition of $G_R$ that $G_R$

is just $\mathbb E R_t$.

We test the condition $\beta \mathbb E R_t < 1$ in the code below.

## Implementation

ζ_draws: jnp.ndarray

def create_ifp(γ=1.5,

β=0.96,

P=np.array([(0.9, 0.1),

(0.1, 0.9)]),

a_r=0.16,

b_r=0.0,

a_y=0.2,

b_y=0.5,

shock_draw_size=100,

grid_max=100,

grid_size=100,

seed=1234):

def create_ifp(

γ=1.5, # Utility parameter

β=0.96, # Discount factor

P=jnp.array([(0.9, 0.1), # Default Markov chain for Z

(0.1, 0.9)]),

a_r=0.16, # Volatility term in R shock

b_r=0.0, # Mean shift R shock

a_y=0.2, # Volatility term in Y shock

b_y=0.5, # Mean shift Y shock

shock_draw_size=100, # For Monte Carlo

grid_max=100, # Exogenous grid max

grid_size=100, # Exogenous grid size

seed=1234 # Random seed

):

"""

Create an instance of IFP with the given parameters.

"""

# Test stability assuming {R_t} is IID and adopts the lognormal

# specification given below. The test is then β E R_t < 1.

# Test stability assuming {R_t} is IID and ln R ~ N(b_r, a_r)

ER = np.exp(b_r + a_r**2 / 2)

assert β * ER < 1, "Stability condition failed."

# Convert to JAX arrays

P = jnp.array(P)

# Generate random draws using JAX

key = jax.random.PRNGKey(seed)

key, subkey1, subkey2 = jax.random.split(key, 3)

subkey1, subkey2 = jax.random.split(key)

η_draws = jax.random.normal(subkey1, (shock_draw_size,))

ζ_draws = jax.random.normal(subkey2, (shock_draw_size,))

s_grid = jnp.linspace(0, grid_max, grid_size)

Here's the Coleman-Reffett operator using JAX:

```{code-cell} ipython3

@jax.jit

def K(ae_vals, c_vals, ifp):

def K(

a_in: jnp.array, # a_in[i, z] is an asset grid

c_in: jnp.array, # c_in[i, z] = consumption at a_in[i, z]

ifp: IFP

):

"""

The Coleman--Reffett operator for the income fluctuation problem,

using the endogenous grid method with JAX.

* ifp is an instance of IFP

* ae_vals[i, z] is an asset grid

* c_vals[i, z] is consumption at ae_vals[i, z]

"""

# Extract parameters from ifp

γ, β, P = ifp.γ, ifp.β, ifp.P

a_r, b_r, a_y, b_y = ifp.a_r, ifp.b_r, ifp.a_y, ifp.b_y

s_grid, η_draws, ζ_draws = ifp.s_grid, ifp.η_draws, ifp.ζ_draws

γ, β, P, a_r, b_r, a_y, b_y, s_grid, η_draws, ζ_draws = ifp

n = len(P)

# Allocate memory

c_out = jnp.empty_like(c_vals)

# Obtain c_i at each s_i, z, store in c_out[i, z], computing

# the expectation term by Monte Carlo

def compute_expectation(s, z):

"""Compute expectation for given s and z"""

def inner_expectation(z_hat):

# Vectorize over shocks

def compute_term(η, ζ):

R_hat = R(z_hat, ζ, a_r, b_r)

Y_hat = Y(z_hat, η, a_y, b_y)

a_val = R_hat * s + Y_hat

# Interpolate consumption

c_interp = jnp.interp(a_val, ae_vals[:, z_hat], c_vals[:, z_hat])

U = u_prime(c_interp, γ)

return R_hat * U

c_interp = jnp.interp(a_val, a_in[:, z_hat], c_in[:, z_hat])

mu = u_prime(c_interp, γ)

return R_hat * mu

# Vectorize over all shock combinations

η_grid, ζ_grid = jnp.meshgrid(η_draws, ζ_draws, indexing='ij')

terms = vmap(vmap(compute_term))(η_grid, ζ_grid)

return P[z, z_hat] * jnp.mean(terms)

# Sum over z_hat states

Ez = jnp.sum(vmap(inner_expectation)(jnp.arange(n)))

return u_prime_inv(β * Ez, γ)

# Vectorize over s_grid and z

c_out = vmap(vmap(compute_expectation, in_axes=(None, 0)),

in_axes=(0, None))(s_grid, jnp.arange(n))

compute_exp_v1 = vmap(compute_expectation, in_axes=(None, 0))

compute_exp_v2 = vmap(compute_exp_v1, in_axes=(0, None))

c_out = compute_exp_v2(s_grid, jnp.arange(n))

# Calculate endogenous asset grid

ae_out = s_grid[:, None] + c_out

# Fixing a consumption-asset pair at (0, 0) improves interpolation

a_out = s_grid[:, None] + c_out

# Fix consumption-asset pair at (0, 0)

c_out = c_out.at[0, :].set(0)

ae_out = ae_out.at[0, :].set(0)

a_out = a_out.at[0, :].set(0)

return ae_out, c_out

return a_out, c_out

```

The next function solves for an approximation of the optimal consumption policy

via time iteration using JAX:

```{code-cell} ipython3

def solve_model_time_iter(

model, # Class with model information

a_vec, # Initial condition for assets

σ_vec, # Initial condition for consumption

tol=1e-4,

max_iter=1000,

verbose=True,

print_skip=25

):

# Set up loop

i = 0

error = tol + 1

while i < max_iter and error > tol:

a_new, σ_new = K(a_vec, σ_vec, model)

error = jnp.max(jnp.abs(σ_vec - σ_new))

@jax.jit

def solve_model(

ifp: IFP,

c_init: jnp.ndarray, # Initial guess of σ on grid endogenous grid

a_init: jnp.ndarray, # Initial endogenous grid

tol: float = 1e-5,

max_iter: int = 1000

) -> jnp.ndarray:

" Solve the model using time iteration with EGM. "

def condition(loop_state):

c_in, a_in, i, error = loop_state

return (error > tol) & (i < max_iter)

def body(loop_state):

c_in, a_in, i, error = loop_state

c_out, a_out = K(c_in, a_in, ifp)

error = jnp.max(jnp.abs(c_out - c_in))

i += 1

if verbose and i % print_skip == 0:

print(f"Error at iteration {i} is {error}.")

a_vec, σ_vec = a_new, σ_new

return c_out, a_out, i, error

if error > tol:

print("Failed to converge!")

elif verbose:

print(f"\nConverged in {i} iterations.")

i, error = 0, tol + 1

initial_state = (c_init, a_init, i, error)

final_loop_state = jax.lax.while_loop(condition, body, initial_state)

c_out, a_out, i, error = final_loop_state

return a_new, σ_new

return c_out, a_out

```

Now we can create an instance and solve the model using JAX:

Let's generate an approximation solution with JAX:

```{code-cell} ipython3

a_star, σ_star = solve_model_time_iter(ifp, a_init, σ_init, print_skip=5)

a_star, σ_star = solve_model(ifp, a_init, σ_init)

```

Let's try it again with a timer.

```{code-cell} python3

with qe.Timer(precision=8):

a_star, σ_star = solve_model(ifp, a_init, σ_init)

a_star.block_until_ready()

```

## Simulation

as representing an optimal policy associated with a given model `ifp`

```{code-cell} ipython3

@jax.jit

def simulate_household(

key, a_0, z_idx_0, c_vec, a_vec, ifp, T

):

Now we write a function to simulate many households in parallel.

```{code-cell} ipython3

@partial(jax.jit, static_argnums=(3, 4, 5))

def compute_asset_stationary(

c_vec, a_vec, ifp, num_households=50_000, T=500, seed=1234

):

)

assets = sim_all_households(keys, a_0_vector, z_idx_0_vector, c_vec, a_vec, ifp, T)

return np.array(assets)

return jnp.array(assets)

```

We'll need some inequality measures for visualization, so let's define them first:

n_z = len(ifp.P)

a_init = s_grid[:, None] * jnp.ones(n_z)

c_init = a_init

a_vec, c_vec = solve_model_time_iter(ifp, a_init, c_init)

a_vec, c_vec = solve_model(ifp, a_init, c_init)

assets = compute_asset_stationary(c_vec, a_vec, ifp, num_households=200_000)

# Compute Gini coefficient for the plot

n_z_temp = len(ifp_temp.P)

a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)

c_init_temp = a_init_temp

a_vec_temp, c_vec_temp = solve_model_time_iter(

ifp_temp, a_init_temp, c_init_temp, verbose=False

a_vec_temp, c_vec_temp = solve_model(

ifp_temp, a_init_temp, c_init_temp

)

# Simulate households

n_z_temp = len(ifp_temp.P)

a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)

c_init_temp = a_init_temp

a_vec_temp, c_vec_temp = solve_model_time_iter(

ifp_temp, a_init_temp, c_init_temp, verbose=False

a_vec_temp, c_vec_temp = solve_model(

ifp_temp, a_init_temp, c_init_temp

)

# Simulate households