Skip to content

Add visual comparison of HPI, VFI and OPI outputs #171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 1, 2024
Merged

Add visual comparison of HPI, VFI and OPI outputs #171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
494a7f1
Add visual comparison of HPI, VFI and OPI outputs using plots.
kp992 Apr 1, 2024
968bd66
Add visual comparison of three outputs
kp992 Apr 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 43 additions & 23 deletions lectures/opt_invest.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,6 @@ def value_function_iteration(model, tol=1e-5):

For OPI we will use a compiled JAX `lax.while_loop` operation to speed execution.


```{code-cell} ipython3
def opi_loop(params, sizes, arrays, m, tol, max_iter):
"""
Expand Down Expand Up @@ -389,7 +388,6 @@ def optimistic_policy_iteration(model, m=10, tol=1e-5, max_iter=10_000):

Here's HPI


```{code-cell} ipython3
def howard_policy_iteration(model, maxiter=250):
"""
Expand All @@ -408,70 +406,92 @@ def howard_policy_iteration(model, maxiter=250):
return σ
```


```{code-cell} ipython3
:tags: [hide-output]

model = create_investment_model()
constants, sizes, arrays = model
β, a_0, a_1, γ, c = constants
y_size, z_size = sizes
y_grid, z_grid, Q = arrays
```

```{code-cell} ipython3
:tags: [hide-output]

print("Starting HPI.")
qe.tic()
out = howard_policy_iteration(model)
σ_star_hpi = howard_policy_iteration(model)
elapsed = qe.toc()
print(out)
print(σ_star_hpi)
print(f"HPI completed in {elapsed} seconds.")
```

Here's the plot of the Howard policy, as a function of $y$ at the highest and lowest values of $z$.

```{code-cell} ipython3
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(y_grid, y_grid, "k--", label="45")
ax.plot(y_grid, y_grid[σ_star_hpi[:, 1]], label="$\\sigma^{*}_{HPI}(\cdot, z_1)$")
ax.plot(y_grid, y_grid[σ_star_hpi[:, -1]], label="$\\sigma^{*}_{HPI}(\cdot, z_N)$")
ax.legend(fontsize=12)
plt.show()
```

```{code-cell} ipython3
:tags: [hide-output]

print("Starting VFI.")
qe.tic()
out = value_function_iteration(model)
σ_star_vfi = value_function_iteration(model)
elapsed = qe.toc()
print(out)
print(σ_star_vfi)
print(f"VFI completed in {elapsed} seconds.")
```

Here's the plot of the VFI, as a function of $y$ at the highest and lowest values of $z$.

```{code-cell} ipython3
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(y_grid, y_grid, "k--", label="45")
ax.plot(y_grid, y_grid[σ_star_vfi[:, 1]], label="$\\sigma^{*}_{VFI}(\cdot, z_1)$")
ax.plot(y_grid, y_grid[σ_star_vfi[:, -1]], label="$\\sigma^{*}_{VFI}(\cdot, z_N)$")
ax.legend(fontsize=12)
plt.show()
```

```{code-cell} ipython3
:tags: [hide-output]

print("Starting OPI.")
qe.tic()
out = optimistic_policy_iteration(model, m=100)
σ_star_opi = optimistic_policy_iteration(model, m=100)
elapsed = qe.toc()
print(out)
print(σ_star_opi)
print(f"OPI completed in {elapsed} seconds.")
```

Here's the plot of the Howard policy, as a function of $y$ at the highest and lowest values of $z$.
Here's the plot of the optimal policy, as a function of $y$ at the highest and lowest values of $z$.

```{code-cell} ipython3
model = create_investment_model()
constants, sizes, arrays = model
β, a_0, a_1, γ, c = constants
y_size, z_size = sizes
y_grid, z_grid, Q = arrays
```

```{code-cell} ipython3
σ_star = howard_policy_iteration(model)

fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(y_grid, y_grid, "k--", label="45")
ax.plot(y_grid, y_grid[σ_star[:, 1]], label="$\\sigma^*(\cdot, z_1)$")
ax.plot(y_grid, y_grid[σ_star[:, -1]], label="$\\sigma^*(\cdot, z_N)$")
ax.plot(y_grid, y_grid[σ_star_opi[:, 1]], label="$\\sigma^{*}_{OPI}(\cdot, z_1)$")
ax.plot(y_grid, y_grid[σ_star_opi[:, -1]], label="$\\sigma^{*}_{OPI}(\cdot, z_N)$")
ax.legend(fontsize=12)
plt.show()
```

We observe that all the solvers produce the same output from the above three plots.


Let's plot the time taken by each of the solvers and compare them.

```{code-cell} ipython3
m_vals = range(5, 600, 40)
```

```{code-cell} ipython3
model = create_investment_model()
print("Running Howard policy iteration.")
qe.tic()
σ_pi = howard_policy_iteration(model)
Expand Down
86 changes: 51 additions & 35 deletions lectures/opt_savings_2.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ kernelspec:
name: python3
---


# Optimal Savings II: Alternative Algorithms

```{include} _admonition/gpu.md
Expand Down Expand Up @@ -65,7 +64,6 @@ We will use the following imports:
import quantecon as qe
import jax
import jax.numpy as jnp
from collections import namedtuple
import matplotlib.pyplot as plt
import time
```
Expand Down Expand Up @@ -171,7 +169,6 @@ def get_greedy(v, params, sizes, arrays):
return jnp.argmax(B(v, params, sizes, arrays), axis=-1)

get_greedy = jax.jit(get_greedy, static_argnums=(2,))

```

We define a function to compute the current rewards $r_\sigma$ given policy $\sigma$,
Expand Down Expand Up @@ -248,7 +245,6 @@ def T_σ(v, σ, params, sizes, arrays):
T_σ = jax.jit(T_σ, static_argnums=(3,))
```


The function below computes the value $v_\sigma$ of following policy $\sigma$.

This lifetime value is a function $v_\sigma$ that satisfies
Expand Down Expand Up @@ -325,11 +321,8 @@ def get_value(σ, params, sizes, arrays):
return jax.scipy.sparse.linalg.bicgstab(partial_L_σ, r_σ)[0]

get_value = jax.jit(get_value, static_argnums=(2,))

```



## Iteration


Expand Down Expand Up @@ -374,7 +367,6 @@ iterate_policy_operator = jax.jit(iterate_policy_operator,
static_argnums=(4,))
```


## Solvers

Now we define the solvers, which implement VFI, HPI and OPI.
Expand All @@ -395,7 +387,6 @@ def value_function_iteration(model, tol=1e-5):

For OPI we will use a compiled JAX `lax.while_loop` operation to speed execution.


```{code-cell} ipython3
def opi_loop(params, sizes, arrays, m, tol, max_iter):
"""
Expand Down Expand Up @@ -436,7 +427,6 @@ def optimistic_policy_iteration(model, m=10, tol=1e-5, max_iter=10_000):
return σ_star
```


Here's HPI.

```{code-cell} ipython3
Expand All @@ -457,9 +447,9 @@ def howard_policy_iteration(model, maxiter=250):
return σ
```

## Plots
## Tests

Create a model for consumption, perform policy iteration, and plot the resulting optimal policy function.
Let's create a model for consumption, and plot the resulting optimal policy function using all the three algorithms and also check the time taken by each solver.

```{code-cell} ipython3
model = create_consumption_model()
Expand All @@ -470,55 +460,82 @@ w_size, y_size = sizes
w_grid, y_grid, Q = arrays
```

```{code-cell} ipython3
print("Starting HPI.")
start_time = time.time()
σ_star_hpi = howard_policy_iteration(model)
elapsed = time.time() - start_time
print(f"HPI completed in {elapsed} seconds.")
```

```{code-cell} ipython3
---
mystnb:
figure:
caption: Optimal policy function
name: optimal-policy-function
caption: Optimal policy function (HPI)
name: optimal-policy-function-hpi
---
σ_star = howard_policy_iteration(model)

fig, ax = plt.subplots()
ax.plot(w_grid, w_grid, "k--", label="45")
ax.plot(w_grid, w_grid[σ_star[:, 1]], label="$\\sigma^*(\cdot, y_1)$")
ax.plot(w_grid, w_grid[σ_star[:, -1]], label="$\\sigma^*(\cdot, y_N)$")
ax.plot(w_grid, w_grid[σ_star_hpi[:, 1]], label="$\\sigma^{*}_{HPI}(\cdot, y_1)$")
ax.plot(w_grid, w_grid[σ_star_hpi[:, -1]], label="$\\sigma^{*}_{HPI}(\cdot, y_N)$")
ax.legend()
plt.show()
```

## Tests

Here's a quick test of the timing of each solver.

```{code-cell} ipython3
model = create_consumption_model()
```

```{code-cell} ipython3
print("Starting HPI.")
print("Starting VFI.")
start_time = time.time()
out = howard_policy_iteration(model)
σ_star_vfi = value_function_iteration(model)
elapsed = time.time() - start_time
print(f"HPI completed in {elapsed} seconds.")
print(f"VFI completed in {elapsed} seconds.")
```

```{code-cell} ipython3
print("Starting VFI.")
start_time = time.time()
out = value_function_iteration(model)
elapsed = time.time() - start_time
print(f"VFI completed in {elapsed} seconds.")
---
mystnb:
figure:
caption: Optimal policy function (VFI)
name: optimal-policy-function-vfi
---

fig, ax = plt.subplots()
ax.plot(w_grid, w_grid, "k--", label="45")
ax.plot(w_grid, w_grid[σ_star_vfi[:, 1]], label="$\\sigma^{*}_{VFI}(\cdot, y_1)$")
ax.plot(w_grid, w_grid[σ_star_vfi[:, -1]], label="$\\sigma^{*}_{VFI}(\cdot, y_N)$")
ax.legend()
plt.show()
```

```{code-cell} ipython3
print("Starting OPI.")
start_time = time.time()
out = optimistic_policy_iteration(model, m=100)
σ_star_opi = optimistic_policy_iteration(model, m=100)
elapsed = time.time() - start_time
print(f"OPI completed in {elapsed} seconds.")
```

```{code-cell} ipython3
---
mystnb:
figure:
caption: Optimal policy function (OPI)
name: optimal-policy-function-opi
---

fig, ax = plt.subplots()
ax.plot(w_grid, w_grid, "k--", label="45")
ax.plot(w_grid, w_grid[σ_star_opi[:, 1]], label="$\\sigma^{*}_{OPI}(\cdot, y_1)$")
ax.plot(w_grid, w_grid[σ_star_opi[:, -1]], label="$\\sigma^{*}_{OPI}(\cdot, y_N)$")
ax.legend()
plt.show()
```

We observe that all the solvers produce the same output from the above three plots.

Now, let's create a plot to visualize the time differences among these algorithms.

```{code-cell} ipython3
def run_algorithm(algorithm, model, **kwargs):
start_time = time.time()
Expand All @@ -530,7 +547,6 @@ def run_algorithm(algorithm, model, **kwargs):
```

```{code-cell} ipython3
model = create_consumption_model()
σ_pi, pi_time = run_algorithm(howard_policy_iteration, model)
σ_vfi, vfi_time = run_algorithm(value_function_iteration, model, tol=1e-5)

Expand Down