QuantEcon / QuantEcon.py / 18963651100

93%

push

github

Optimize draw (#811)

The key optimization replaces a Python loop with vectorized NumPy operations in the `draw` function's multiple sample case.

**What changed:**
- Replaced the explicit Python loop `for i in range(size): out[i] = searchsorted(cdf, rs[i])` with a single vectorized call: `out = np.searchsorted(cdf, rs, side='right')`
- Removed the separate `np.empty` allocation since `np.searchsorted` returns the output array directly

**Why this is faster:**
The original code performs `size` individual calls to the custom `searchsorted` function in Python, each requiring loop overhead and function call overhead. The optimized version leverages NumPy's highly optimized C implementation that processes the entire array in one operation, eliminating Python loop overhead entirely.

**Performance characteristics:**
- Massive speedups for large sample sizes (857% faster for 1000 samples, 934% for 500 samples)  
- Modest improvements for small sample sizes (35-40% faster for 10-100 samples)
- Single draws remain unchanged, preserving the custom implementation's behavior
- Edge cases like `size=0` show slight regression due to NumPy's overhead for empty arrays, but these are uncommon scenarios

The optimization is most effective when `size` is an integer (vectorizable case), while preserving the original behavior for single draws and non-integer sizes.

Co-authored-by: codeflash-ai[bot] <148906541+codeflash-ai[bot]@users.noreply.github.com>

1 of 1 new or added line in 1 file covered. (100.0%)

7708 of 8326 relevant lines covered (92.58%)

2.78 hits per line