Skip to content

v4.2.0: Numba JIT compilation for TKE closure

Latest

Choose a tag to compare

@samfrederick samfrederick released this 10 Feb 02:27

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_length and compute_tke_closure_column) contained Python-level for k loops over vertical levels
  • After JIT compilation: ~200x speedup per column call (from ~1 ms to ~5 μs)

Changes

  • Numba @njit on 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_end extended from 3600s (1 hour) to 7200s (2 hours)
  • requirements.txt — Added numba

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.