PolicyEngine-UK compatibility: Loosen SS tolerances and add Numba optimizations

[vahid-ahmadi]

Performance Optimizations for Steady State Solver

Summary

• Loosened convergence tolerances for faster development/exploratory runs
• Added optional Numba JIT compilation for compute-intensive functions

Changes

1. Solver Tolerance Adjustments (default_parameters.json)

Parameter	Before	After	Impact
mindist_SS	1e-9	1e-3	GE loop converges faster
RC_SS	1e-8	1e-4	Resource constraint check passes earlier

Benchmark results:

• Before: ~15 min runtime
• After: ~7 min runtime (~55% faster)
• Result accuracy: < 0.01% difference in key outputs (Y, C, K, L, r, w)

2. Optional Numba JIT Compilation

Added Numba-optimized versions of hot-path functions:

• household.py: marg_ut_cons, marg_ut_labor
• tax.py: ETR_wealth, MTR_wealth
• txfunc.py: DEP tax rate calculations

Features:

• Falls back gracefully to NumPy if Numba not installed
• Uses @njit(cache=True) for compilation caching
• Parallelized with prange where applicable

Test Plan

• Run OG-UK example simulation
• Verify results match within acceptable tolerance
• Confirm backward compatibility (works without Numba)

🤖 Generated with Claude Code

[vahid-ahmadi]

@MaxGhenis @nikhilwoodruff

[nikhilwoodruff]

These are parameter changes right? Why do we need to change the defaults, we can just change them in og-uk no?

[vahid-ahmadi]

These are parameter changes right? Why do we need to change the defaults, we can just change them in og-uk no?

I started working on Numba JIT optimizations for hot-path functions (compiling Python loops to fast machine code). Still, after benchmarking, most of the speedup is actually coming from the changes to the tolerance parameters (mindist_SS, RC_SS).

I can move those parameter changes to OG-UK via p.update_specifications() instead of changing OG-Core defaults.

[jdebacker]

@vahid-ahmadi I echo @nikhilwoodruff, let's not change the default values in OG-Core. These parameters can be adjusted for particular runs (e.g., you can set a looser tolerance if you want a quicker solution).

[jdebacker]

This max value seems too high to allow. For example, module units for labor supply are typically less than 0.5. So this would mean an error of 20% or more is ok. I think that's too much for any simulation.

[jdebacker]

@vahid-ahmadi Thanks for thinking about improvements to the solution of OG-Core. Can you report some diffs of the Numba changes only (i.e., keep the default tolerances and run with only the numba changes). What is the speed up? These are pretty small loops (and sometimes vectorized operations) so I'm skeptical there would be much improvement.