Numba JIT compilation for the Deardorff TKE closure, eliminating the dominant performance bottleneck.
Performance
- Profiling showed the TKE closure consumed 86% of total runtime in both 2D and 3D simulations
- The inner column functions (
_compute_mixing_lengthandcompute_tke_closure_column) contained Python-levelfor kloops over vertical levels - After JIT compilation: ~200x speedup per column call (from ~1 ms to ~5 μs)
Changes
- Numba
@njiton TKE closure column functions —_compute_mixing_length_jit()and_compute_tke_closure_column_jit()accept only NumPy arrays and scalar primitives; Python wrappers preserve the original API for backward compatibility config/cbl_3d.yaml—t_endextended from 3600s (1 hour) to 7200s (2 hours)requirements.txt— Addednumba
Backward Compatibility
- All public APIs unchanged — callers do not need any modifications
- First call incurs a one-time JIT compilation overhead (~1-2s); subsequent calls run at compiled speed
- Numerical results identical to v4.1.0
See RELEASE_NOTES.md for full details.