promote: refresh 14 pool lectures whose canonical source moved

lectures/ifp_advanced.md

@@ -32,7 +32,7 @@ In addition to what's in Anaconda, this lecture will need the following librarie
 ---
 tags: [hide-output]
 ---
-!pip install quantecon
+!pip install quantecon jax
 ```
 
 ## Overview
@@ -356,7 +356,7 @@ def create_ifp(
     assert β * ER < 1, "Stability condition failed."
 
     # Generate random draws using JAX
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     subkey1, subkey2 = jax.random.split(key)
     η_draws = jax.random.normal(subkey1, (shock_draw_size,))
     ζ_draws = jax.random.normal(subkey2, (shock_draw_size,))
@@ -388,8 +388,8 @@ Here's the Coleman-Reffett operator using JAX:
 
 ```{code-cell} ipython3
 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]
+        a_in: jnp.array,   # a_in[i, z] is an asset grid
         ifp: IFP
     ):
     """
@@ -430,7 +430,7 @@ def K(
     c_out = c_out.at[0, :].set(0)
     a_out = a_out.at[0, :].set(0)
 
-    return a_out, c_out
+    return c_out, a_out
 ```
 
 The next function solves for an approximation of the optimal consumption policy
@@ -487,15 +487,15 @@ a_init = σ_init.copy()
 Let's generate an approximation solution with JAX:
 
 ```{code-cell} ipython3
-a_star, σ_star = solve_model(ifp, a_init, σ_init)
+σ_star, a_star = solve_model(ifp, σ_init, a_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()
+    σ_star, a_star = solve_model(ifp, σ_init, a_init)
+    σ_star.block_until_ready()
 ```
 
 ## Simulation
@@ -586,7 +586,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(
@@ -642,7 +642,7 @@ s_grid = ifp.s_grid
 n_z = len(ifp.P)
 a_init = s_grid[:, None] * jnp.ones(n_z)
 c_init = a_init
-a_vec, c_vec = solve_model(ifp, a_init, c_init)
+c_vec, a_vec = solve_model(ifp, c_init, a_init)
 assets = compute_asset_stationary(c_vec, a_vec, ifp, num_households=200_000)
 
 # Compute Gini coefficient for the plot
@@ -734,8 +734,8 @@ for a_r in a_r_vals:
     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(
-        ifp_temp, a_init_temp, c_init_temp
+    c_vec_temp, a_vec_temp = solve_model(
+        ifp_temp, c_init_temp, a_init_temp
     )
 
     # Simulate households
@@ -811,8 +811,8 @@ for a_y in a_y_vals:
     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(
-        ifp_temp, a_init_temp, c_init_temp
+    c_vec_temp, a_vec_temp = solve_model(
+        ifp_temp, c_init_temp, a_init_temp
     )
 
     # Simulate households

lectures/ifp_egm.md

@@ -890,7 +890,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(

lectures/ifp_egm_transient_shocks.md

@@ -411,7 +411,7 @@ def create_ifp(r=0.01,
                shock_draw_size=100,
                seed=1234):
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     s = jnp.linspace(0, savings_grid_max, savings_grid_size)
     Π, z_grid = jnp.array(Π), jnp.array(z_grid)
     R = 1 + r
@@ -779,7 +779,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(

lectures/inventory_q.md

@@ -355,13 +355,12 @@ At each step, we draw a demand shock from the geometric distribution and update
 
 ```{code-cell} ipython3
 @numba.jit(nopython=True)
-def sim_inventories(ts_length, σ, p, X_init=0, seed=0):
+def sim_inventories(ts_length, σ, p, rng, X_init=0):
     """Simulate inventory dynamics under policy σ."""
-    np.random.seed(seed)
     X = np.zeros(ts_length, dtype=np.int32)
     X[0] = X_init
     for t in range(ts_length - 1):
-        d = np.random.geometric(p) - 1
+        d = rng.geometric(p) - 1
         X[t+1] = max(X[t] - d, 0) + σ[X[t]]
     return X
 ```
@@ -373,8 +372,8 @@ a large order to replenish stock (the upward jumps), after which inventory
 gradually declines as demand is served.
 
 ```{code-cell} ipython3
-def plot_ts(ts_length=200, fontsize=10):
-    X = sim_inventories(ts_length, σ_star, p)
+def plot_ts(ts_length=200, fontsize=10, seed=0):
+    X = sim_inventories(ts_length, σ_star, p, np.random.default_rng(seed))
     fig, ax = plt.subplots()
 
     ax.plot(X, label=r"$X_t$", alpha=0.7)
@@ -588,8 +587,7 @@ At specified step counts (given by `snapshot_steps`), we record the current gree
 ```{code-cell} ipython3
 @numba.jit(nopython=True)
 def q_learning_kernel(K, p, c, κ, β, n_steps, X_init,
-                      ε_init, ε_min, ε_decay, q_init, snapshot_steps, seed):
-    np.random.seed(seed)
+                      ε_init, ε_min, ε_decay, q_init, snapshot_steps, rng):
     q = np.full((K + 1, K + 1), q_init)
     n = np.zeros((K + 1, K + 1))       # visit counts for learning rate
     ε = ε_init
@@ -600,7 +598,7 @@ def q_learning_kernel(K, p, c, κ, β, n_steps, X_init,
 
     # Initialize state and action
     x = X_init
-    a = np.random.randint(0, K - x + 1)
+    a = rng.integers(0, K - x + 1)
 
     for t in range(n_steps):
         # Record policy snapshot if needed
@@ -609,7 +607,7 @@ def q_learning_kernel(K, p, c, κ, β, n_steps, X_init,
             snap_idx += 1
 
         # === Draw D_{t+1} and observe outcome ===
-        d = np.random.geometric(p) - 1
+        d = rng.geometric(p) - 1
         reward = min(x, d) - c * a - κ * (a > 0)
         x_next = max(x - d, 0) + a
 
@@ -629,8 +627,8 @@ def q_learning_kernel(K, p, c, κ, β, n_steps, X_init,
 
         # === Behavior policy: ε-greedy (uses a_next, the argmax action) ===
         x = x_next
-        if np.random.random() < ε:
-            a = np.random.randint(0, K - x + 1)
+        if rng.random() < ε:
+            a = rng.integers(0, K - x + 1)
         else:
             a = a_next
         ε = max(ε_min, ε * ε_decay)
@@ -648,8 +646,9 @@ def q_learning(model, n_steps=20_000_000, X_init=0,
     K = len(x_values) - 1
     if snapshot_steps is None:
         snapshot_steps = np.array([], dtype=np.int64)
+    rng = np.random.default_rng(seed)
     return q_learning_kernel(K, p, c, κ, β, n_steps, X_init,
-                             ε_init, ε_min, ε_decay, q_init, snapshot_steps, seed)
+                             ε_init, ε_min, ε_decay, q_init, snapshot_steps, rng)
 ```
 
 Next we run $n$ = 5 million steps and take policy snapshots at steps 10,000, 1,000,000, and $n$.
@@ -726,7 +725,8 @@ X_init = K // 2
 sim_seed = 5678
 
 # Optimal policy
-X_opt = sim_inventories(ts_length, σ_star, p, X_init, seed=sim_seed)
+X_opt = sim_inventories(ts_length, σ_star, p,
+                        np.random.default_rng(sim_seed), X_init)
 axes[0].plot(X_opt, alpha=0.7)
 axes[0].set_ylabel("inventory")
 axes[0].set_title("Optimal (VFI)")
@@ -735,7 +735,8 @@ axes[0].set_ylim(0, K + 2)
 # Q-learning snapshots
 for i in range(n_snaps):
     σ_snap = snapshots[i]
-    X = sim_inventories(ts_length, σ_snap, p, X_init, seed=sim_seed)
+    X = sim_inventories(ts_length, σ_snap, p,
+                        np.random.default_rng(sim_seed), X_init)
     axes[i + 1].plot(X, alpha=0.7)
     axes[i + 1].set_ylabel("inventory")
     axes[i + 1].set_title(f"Step {snap_steps[i]:,}")