@@ -101,13 +101,13 @@ threshold w*.
101101
102102In addition to what's in Anaconda, this lecture will need the QE library:
103103
104- ``` python
104+ ``` {code-cell} ipython3
105105#!pip install quantecon # Uncomment if necessary
106106```
107107
108108We use the following imports:
109109
110- ``` python
110+ ``` {code-cell} ipython3
111111from quantecon.markov import tauchen
112112import jax.numpy as jnp
113113import jax
@@ -119,7 +119,7 @@ from functools import partial
119119
120120First, we implement the successive approximation algorithm:
121121
122- ``` python
122+ ``` {code-cell} ipython3
123123@partial(jit, static_argnums=(0,))
124124def successive_approx(
125125 T, # Operator (callable) - marked as static
@@ -148,7 +148,7 @@ def successive_approx(
148148
149149Let's set up a ` Model ` class to store information needed to solve the model:
150150
151- ``` python
151+ ``` {code-cell} ipython3
152152class Model(NamedTuple):
153153 n: int
154154 w_vals: jnp.ndarray
@@ -160,7 +160,7 @@ class Model(NamedTuple):
160160
161161The function below holds default values and creates a ` Model ` instance:
162162
163- ``` python
163+ ``` {code-cell} ipython3
164164def create_js_with_sep_model(
165165 n: int = 200, # wage grid size
166166 ρ: float = 0.9, # wage persistence
@@ -177,7 +177,7 @@ def create_js_with_sep_model(
177177
178178Here's the Bellman operator for the unemployed worker's value function:
179179
180- ``` python
180+ ``` {code-cell} ipython3
181181@jit
182182def T(v: jnp.ndarray, model: Model) -> jnp.ndarray:
183183 """The Bellman operator for the value of being unemployed."""
@@ -191,7 +191,7 @@ def T(v: jnp.ndarray, model: Model) -> jnp.ndarray:
191191The next function computes the optimal policy under the assumption that v is
192192the value function:
193193
194- ``` python
194+ ``` {code-cell} ipython3
195195@jit
196196def get_greedy(v: jnp.ndarray, model: Model) -> jnp.ndarray:
197197 """Get a v-greedy policy."""
@@ -205,7 +205,7 @@ def get_greedy(v: jnp.ndarray, model: Model) -> jnp.ndarray:
205205
206206Here's a routine for value function iteration:
207207
208- ``` python
208+ ``` {code-cell} ipython3
209209def vfi(model: Model):
210210 """Solve by VFI."""
211211 v_init = jnp.zeros(model.w_vals.shape)
@@ -240,7 +240,7 @@ def get_reservation_wage(σ: jnp.ndarray, model: Model) -> float:
240240
241241Let's solve the model and plot the results:
242242
243- ``` python
243+ ``` {code-cell} ipython3
244244model = create_js_with_sep_model()
245245n, w_vals, P, β, c, α = model
246246v_star, σ_star = vfi(model)
@@ -266,7 +266,7 @@ plt.show()
266266
267267Let's examine how reservation wages change with the separation rate α:
268268
269- ``` python
269+ ``` {code-cell} ipython3
270270α_vals: jnp.ndarray = jnp.linspace(0.0, 1.0, 10)
271271
272272w_star_vec = jnp.empty_like(α_vals)
@@ -289,7 +289,7 @@ plt.show()
289289
290290Now let's simulate the employment dynamics of a single agent under the optimal policy:
291291
292- ``` python
292+ ``` {code-cell} ipython3
293293@jit
294294def weighted_choice(key, probs):
295295 """JAX-compatible weighted random choice."""
@@ -324,7 +324,7 @@ def update_agent(key, is_employed, wage_idx, model, σ_star):
324324 return final_employment, final_wage
325325```
326326
327- ``` python
327+ ``` {code-cell} ipython3
328328def simulate_employment_path(
329329 model: Model, # Model details
330330 T: int = 2_000, # Simulation length
@@ -360,7 +360,7 @@ def simulate_employment_path(
360360
361361Let's create a comprehensive plot of the employment simulation:
362362
363- ``` python
363+ ``` {code-cell} ipython3
364364model = create_js_with_sep_model()
365365
366366# Calculate reservation wage for plotting
@@ -426,14 +426,14 @@ The model successfully captures the essential features of labor market dynamics
426426
427427Now let's simulate many agents simultaneously to examine the cross-sectional unemployment rate:
428428
429- ``` python
429+ ``` {code-cell} ipython3
430430# Create vectorized version of update_agent
431431update_agents_vmap = jax.vmap(
432432 update_agent, in_axes=(0, 0, 0, None, None)
433433)
434434```
435435
436- ``` python
436+ ``` {code-cell} ipython3
437437@partial(jit, static_argnums=(3, 4))
438438def _simulate_cross_section_compiled(
439439 key: jnp.ndarray,
@@ -515,7 +515,7 @@ def simulate_cross_section(
515515 return unemployment_rates, employment_matrix
516516```
517517
518- ``` python
518+ ``` {code-cell} ipython3
519519def plot_cross_sectional_unemployment(model: Model):
520520 """
521521 Generate cross-sectional unemployment rate plot for a given model.
@@ -555,7 +555,7 @@ def plot_cross_sectional_unemployment(model: Model):
555555 plt.show()
556556```
557557
558- ``` python
558+ ``` {code-cell} ipython3
559559model = create_js_with_sep_model()
560560plot_cross_sectional_unemployment(model)
561561```
@@ -564,7 +564,7 @@ plot_cross_sectional_unemployment(model)
564564
565565Let's examine how the cross-sectional unemployment rate changes with lower unemployment compensation:
566566
567- ``` python
567+ ``` {code-cell} ipython3
568568model_low_c = create_js_with_sep_model(c=0.5)
569569plot_cross_sectional_unemployment(model_low_c)
570570```
@@ -574,14 +574,14 @@ plot_cross_sectional_unemployment(model_low_c)
574574Create a plot that shows how the steady state cross-sectional unemployment rate
575575changes with unemployment compensation.
576576
577- ``` python
577+ ``` {code-cell} ipython3
578578for _ in range(20):
579579 print('Solution below!')
580580```
581581
582582## Solution
583583
584- ``` python
584+ ``` {code-cell} ipython3
585585c_values = 1.0, 0.8, 0.6, 0.4, 0.2
586586rates = []
587587for c in c_values:
@@ -598,6 +598,6 @@ ax.legend(frameon=False)
598598plt.show()
599599```
600600
601- ``` python
601+ ``` {code-cell} ipython3
602602
603603```
0 commit comments