Commit 670d7c0
Simplify McCall model: Compute reservation wage directly from value function (#695)
* Simplify McCall model: Compute reservation wage directly from value function
This commit simplifies the job search model by eliminating the intermediate policy array and computing the reservation wage directly from the value function.
Key changes:
- Removed `get_greedy` function entirely
- Modified `get_reservation_wage` to compute directly from value function v_u
- Updated `update_agent` to use reservation wage (scalar) instead of policy array
- Updated all simulation functions to work with reservation wage
Benefits:
- Simpler, more intuitive code
- More efficient (no need to store full policy array)
- Clearer conceptually (directly compute and use the threshold)
The reservation wage is now computed by finding the lowest wage where:
accept value >= reject value, which directly implements the optimal policy.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
* misc
* Remove unnecessary @jax.jit decorator from update_agent
Remove @jax.jit from the update_agent function. 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 prevent optimizations.
Performance benchmarks show this change provides 0-16% speedup for typical
problem sizes, as the top-level JIT compilation traces through the entire
computation graph more efficiently without the nested decorator.
This follows JAX best practices: only JIT-compile top-level functions and
let them trace through intermediate functions naturally.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
---------
Co-authored-by: Claude <noreply@anthropic.com>1 parent 5957e45 commit 670d7c0
1 file changed
+172
-108
lines changed
0 commit comments