Improve McCall model lecture: Fix beta unpacking and enhance readability

[jstac]

Summary

This PR improves the McCall job search model lecture (mccall_model.md) with several enhancements focused on code clarity and pedagogical value.

Key improvements:

1. Fixed beta unpacking in Comparative statics section (lines 506-516, 523)

   • Now unpacks model parameters (c, β, w, q = model) before use
   • Consistent with patterns used throughout the lecture
   • More readable and follows JAX functional programming style

2. Enhanced derivation of nonlinear equation in h (lines 625-672)

   • Added step-by-step algebraic transformations under "Take 2" section
   • Shows explicitly: multiply by q(w'), sum over w', multiply by β, add c, apply definition
   • Makes the mathematical derivation much clearer for students

3. Improved JAX solution in Exercise 1 (lines 817-847)

   • Extracted wage drawing logic into documented draw_wage() function
   • Added clear explanation of inverse transform method
   • More modular and easier to understand for students learning JAX

4. Broke vectorization into clear steps (lines 887-888, 1028-1034)

   • Separated vmap calls with explanatory comments
   • Shows progression: create vectorized function, then apply it
   • Better demonstrates JAX's composition of vectorization operations

5. Simplified Exercise 2 solution (lines 985, 990, 1006-1007, 1028-1034)

   • Removed redundant keyword arguments (c=c, β=β → c, β)
   • Removed unnecessary meshgrid (vmap handles broadcasting)
   • Unpacked model directly (model instead of model.c, model.β, ...)
   • Made while_loop pattern clearer with explicit final_state variable

Testing

• ✅ Converted to Python script using jupytext
• ✅ Verified script runs successfully
• ✅ All code examples maintain functionality
• ✅ Converted back to markdown format

Benefits

• More consistent code style throughout the lecture
• Better pedagogical progression for students learning JAX
• Clearer mathematical derivations
• More readable and maintainable code examples

🤖 Generated with Claude Code

[jstac]

Detailed Change Breakdown

1. Beta Unpacking (Comparative statics section)
Before:

model_default = McCallModel()
v_init = model_default.w / (1 - model_default.β)

After:

model_default = McCallModel()
c, β, w, q = model_default
v_init = w / (1 - β)

This follows the pattern established earlier in the lecture and is more consistent with JAX's functional style.

[jstac]

Mathematical Derivation Enhancement

2. Derivation of nonlinear equation in h (Take 2 section)

The previous version just stated "Substituting this last equation into {eq}`j1` gives..." without showing the steps.

Now we explicitly show:

1. Multiply both sides by q(w')
2. Sum both sides over w' ∈ W
3. Multiply both sides by β
4. Add c to both sides
5. Recognize the left-hand side as h using the definition

This makes the mathematical derivation transparent and easier for students to follow and verify independently.

[jstac]

Pedagogical Improvements in JAX Code

3. Extracting wage drawing logic

Added a documented draw_wage() function that explains the inverse transform method:

def draw_wage(uniform_rv):
    """
    Draw a wage from the distribution q_default using the inverse transform method.
    ...
    """
    return w_default[jnp.searchsorted(cdf, uniform_rv)]

This makes the code more modular and helps students understand the sampling technique.

4. Breaking up vectorization

Changed from:

stop_times = jax.vmap(compute_stop_time_for_c)(c_vals)

To:

compute_stop_time_vectorized = jax.vmap(compute_stop_time_for_c)
stop_times = compute_stop_time_vectorized(c_vals)

Shows clearly that vmap creates a new vectorized function, then we apply it.

[jstac]

Code Simplification in Exercise 2

5. Multiple improvements for clarity

Removed redundant keyword arguments:

# Before: return McCallModelContinuous(c=c, β=β, σ=σ, μ=μ, w_draws=w_draws)
# After:  return McCallModelContinuous(c, β, σ, μ, w_draws)

Simplified model unpacking:

# Before: c, β, σ, μ, w_draws = model.c, model.β, model.σ, model.μ, model.w_draws
# After:  c, β, σ, μ, w_draws = model

Made while_loop clearer:

# Before: h_final, _, _ = jax.lax.while_loop(cond, update, initial_state)
# After:  
final_state = jax.lax.while_loop(cond, update, initial_state)
h_final, _, _ = final_state

Removed unnecessary meshgrid and broke double vmap into two clear steps with comments explaining each vectorization operation.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-694--sunny-cactus-210e3e.netlify.app (2a068dc)

📚 Changed Lecture Pages: mccall_model