Commit 8fb1683
Improve McCall model with separation: Fix function signatures and optimize sampling
This commit makes several improvements to the job search model with separation and Markov wages:
**Function signature fixes:**
- Modified `simulate_employment_path` to accept policy `σ` as a parameter instead of computing it internally
- Updated `update_agent` to use generic parameter name `σ` instead of `σ_star` for better flexibility
- This makes the simulation functions more modular and allows policy reuse across multiple simulations
**Performance optimization:**
- Added `P_cumsum` (precomputed cumulative sum of transition matrix) to the Model class
- Eliminated redundant cumsum computations during Markov chain simulation
- Replaced `weighted_choice` function with direct `jnp.searchsorted` on precomputed cumulative sums
- For n=200 wage states and 100k agents over 200 periods, this eliminates ~20 million O(n) operations
**Documentation:**
- Added explanation of the inverse transform sampling method
- Documented the performance benefits of precomputing cumulative sums
- Clarified the role of each model component
These changes significantly improve simulation performance while making the code more maintainable and reusable.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <[email protected]>1 parent 5e25feb commit 8fb1683
1 file changed
+158
-89
lines changed
0 commit comments