Simplify McCall model: Compute reservation wage directly from value function

[jstac]

Summary

This PR simplifies the job search model with separation and Markov wages by eliminating unnecessary intermediate computations and directly computing the reservation wage from the value function.

Motivation

The original code computed an optimal policy array using get_greedy, then extracted the reservation wage from that policy. Since the optimal policy is a reservation wage strategy (accept all wages above a threshold), we can compute this threshold directly without the intermediate step.

Key changes

1. Removed get_greedy function

• No longer needed since we don't compute a full policy array
• Simplifies the codebase

2. Modified get_reservation_wage function

• Before: Took a policy array σ and extracted the reservation wage
• After: Takes the value function v_u directly and computes the reservation wage by finding where acceptance becomes optimal
• Implementation: Finds lowest wage where accept >= reject

3. Updated update_agent function

• Before: Took policy array σ and looked up σ[wage_idx] to make decisions
• After: Takes reservation wage w_star (scalar) and compares w_vals[wage_idx] >= w_star
• More efficient (no array indexing) and conceptually clearer

4. Updated all simulation functions

• simulate_employment_path: Now takes w_star instead of σ
• _simulate_cross_section_compiled: Now takes w_star instead of σ
• plot_cross_sectional_unemployment: Computes w_star directly
• All code sections that solved the model simplified from:

  v_star = vfi(model)
  σ_star = get_greedy(v_star, model)
  w_star = get_reservation_wage(σ_star, model)

  to:

  v_star = vfi(model)
  w_star = get_reservation_wage(v_star, model)

Benefits

1. Simpler code: Fewer functions, clearer logic flow
2. More efficient:
   • No need to allocate/store full policy array
   • Direct scalar comparisons instead of array lookups in simulations
3. Conceptually clearer: Directly computes and uses the reservation wage threshold
4. Easier to understand: The reservation wage is the fundamental object of interest

Testing

• Converted to Python script and verified all code runs successfully
• Output matches expected behavior (unemployment rates computed correctly)
• All plots and simulations work as before

🤖 Generated with Claude Code

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-695--sunny-cactus-210e3e.netlify.app (fd20ea1)

📚 Changed Lecture Pages: mccall_model_with_sep_markov

[jstac]

Performance optimization: Removed nested @jax.jit decorator

I've removed the @jax.jit decorator from the update_agent function (line 417).

Why this change?

In JAX, when a top-level function is already JIT-compiled (like _simulate_cross_section_compiled), adding @jax.jit to intermediate functions creates nested compilation boundaries that can actually prevent optimizations.

Performance impact

I ran benchmarks comparing both versions:

Problem Size	With nested JIT	Without nested JIT	Speedup
Small (10k agents, 100 steps)	0.159s	0.169s	0.94x
Medium (50k agents, 200 steps)	1.66s	1.43s	1.16x ✓
Large (100k agents, 300 steps)	4.53s	4.30s	1.06x ✓

For typical problem sizes (medium/large), removing the nested JIT provides 5-16% speedup.

JAX best practice

This follows the JAX pattern: only JIT-compile top-level functions and let them trace through intermediate functions naturally. The top-level JIT compilation traces through the entire computation graph and optimizes it holistically.

The code is now simpler and faster! 🚀

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-695--sunny-cactus-210e3e.netlify.app (34a3c62)

📚 Changed Lecture Pages: mccall_model_with_sep_markov