⚡️ Speed up function _make_multidim_func by 46%

[codeflash-ai]

📄 46% (0.46x) speedup for _make_multidim_func in quantecon/quad.py

⏱️ Runtime : 2.94 milliseconds → 2.01 milliseconds (best of 5 runs)

📝 Explanation and details

The optimization achieves a 46% speedup by replacing pure Python/NumPy operations with Numba JIT-compiled functions for the performance-critical ckron and gridmake operations.

Key Optimizations

1. Numba JIT Compilation with Manual Loop Implementation

The optimized code replaces NumPy's high-level vectorized operations (np.tile, np.repeat, np.column_stack) with explicit loops compiled by Numba:

• _gridmake2_numba: Manually constructs the Cartesian product grid using nested loops instead of NumPy's array manipulation functions
• ckron_numba: Implements Kronecker product reduction using an explicit loop instead of functools.reduce
• Decorator settings: @njit(cache=True, fastmath=True) enables aggressive optimizations and compilation caching

2. Why This Is Faster

Eliminating NumPy overhead: While NumPy operations are optimized C code, they still incur overhead from:

• Temporary array allocations (e.g., np.tile, np.repeat create intermediate arrays)
• Python function call overhead in reduce()
• Memory allocation/deallocation cycles

Numba advantages:

• Compiles to native machine code without Python interpreter overhead
• Eliminates intermediate array allocations by writing directly to output arrays
• fastmath=True allows aggressive floating-point optimizations
• cache=True avoids recompilation on subsequent runs

3. Impact on Workloads

The function_references show that _make_multidim_func is called by quadrature functions (qnwcheb, qnwlege, qnwsimp, qnwtrap, qnwbeta, qnwgamma). These are numerical integration routines that:

• Are likely called repeatedly in computational economics/finance simulations
• Build multi-dimensional grids for numerical quadrature
• Benefit significantly from the optimization since grid construction is in the hot path

4. Test Results Analysis

The annotated tests show significant speedups for multi-dimensional cases:

• 2D cases: 48-61% faster (e.g., test_2d_different_args: 60.8% faster)
• 3D cases: 67-83% faster (e.g., test_large_3d: 66.8% faster, test_3d_basic: 82.6% faster)
• High-dimensional: 82% faster (e.g., test_large_high_dim with 4D: 81.7% faster)
• Scalar/1D cases: Minimal improvement or slight regression (JIT compilation overhead dominates for small inputs)

The optimization is most effective when:

• Multiple dimensions are involved (d ≥ 2)
• Grid sizes are moderate to large (total points ≥ 100)
• The function is called repeatedly (JIT compilation cost amortized)

5. Edge Case Handling

The optimized code adds a check for empty arrays (if len(arrays) == 0) to preserve the original behavior of raising TypeError when ckron() is called with no arguments, ensuring backward compatibility.

Summary: This optimization transforms Python/NumPy array operations into compiled native code, eliminating overhead and temporary allocations. It's particularly effective for the multi-dimensional quadrature use case shown in function_references, where grid construction happens frequently during numerical integration.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 20 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import numpy as np
# imports
import pytest
from quantecon.quad import _make_multidim_func

# --- Helper 1D quadrature function for testing ---
def dummy_1d_func(n, a, b):
    """
    Returns n evenly spaced nodes in [a, b] and uniform weights summing to 1.
    """
    n = int(n)
    nodes = np.linspace(a, b, n)
    weights = np.full(n, 1/n)
    return nodes, weights

# --- Basic Test Cases ---

def test_scalar_inputs_basic():
    # All scalar inputs: should just call one_d_func directly
    n, a, b = 3, 0, 1
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 114μs -> 108μs (5.69% faster)
    exp_nodes, exp_weights = dummy_1d_func(n, a, b)

def test_2d_equal_args():
    # 2d, same n, a, b for both dims
    n = [2, 2]
    a = [0, 0]
    b = [1, 1]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 205μs -> 138μs (48.6% faster)
    # Should be grid of [0,1] x [0,1]
    expected_nodes = np.array([[0,0],[1,0],[0,1],[1,1]])
    expected_weights = np.full(4, 0.25)

def test_2d_different_args():
    # 2d, different n, a, b for each dim
    n = [2, 3]
    a = [0, 10]
    b = [1, 13]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 184μs -> 114μs (60.8% faster)
    # First dim: [0,1], second dim: [10,11.5,13]
    exp_nodes_0 = np.linspace(0,1,2)
    exp_nodes_1 = np.linspace(10,13,3)
    # Cartesian product
    expected_nodes = np.array([[0,10],[1,10],[0,11.5],[1,11.5],[0,13],[1,13]])
    # Weights: kron([0.5,0.5], [1/3,1/3,1/3]) = [1/6]*6
    expected_weights = np.full(6, 1/6)

def test_3d_basic():
    # 3d, n, a, b
    n = [2, 2, 2]
    a = [0, 1, 2]
    b = [1, 2, 3]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 223μs -> 122μs (82.6% faster)
    # Each dim: [0,1], [1,2], [2,3]
    grid = np.array(np.meshgrid([0,1],[1,2],[2,3], indexing='ij')).reshape(3,-1).T

def test_args_broadcasting():
    # n is [3,2], a and b are scalars: should be broadcast to [a,a], [b,b]
    n = [3,2]
    a = 0
    b = 1
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 168μs -> 114μs (47.3% faster)
    # Should match calling with a=[0,0], b=[1,1]
    nodes2, weights2 = _make_multidim_func(dummy_1d_func, n, [0,0], [1,1]) # 105μs -> 71.7μs (46.4% faster)

# --- Edge Test Cases ---

def test_single_dim_array():
    # n is array of length 1: should call one_d_func with n[0]
    n = np.array([5])
    a = [0]
    b = [10]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 88.3μs -> 96.9μs (8.95% slower)
    exp_nodes, exp_weights = dummy_1d_func(5, 0, 10)

def test_mixed_scalar_and_array_args():
    # n is [2,3], a is scalar, b is [1,2]
    n = [2,3]
    a = 0
    b = [1,2]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 166μs -> 109μs (52.2% faster)
    # Should match a=[0,0], b=[1,2]
    nodes2, weights2 = _make_multidim_func(dummy_1d_func, n, [0,0], [1,2]) # 106μs -> 69.5μs (52.6% faster)

def test_zero_nodes():
    # n=0 should raise or return empty arrays
    n = 0
    a = 0
    b = 1
    with pytest.raises(ZeroDivisionError):
        _make_multidim_func(dummy_1d_func, n, a, b) # 49.8μs -> 49.6μs (0.464% faster)

def test_non_integer_n():
    # n is float, should be cast to int in dummy_1d_func
    n = [2.0, 3.0]
    a = [0,0]
    b = [1,1]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 166μs -> 105μs (57.2% faster)

def test_negative_n():
    # n negative should error in dummy_1d_func
    n = [-2, 3]
    a = [0,0]
    b = [1,1]
    with pytest.raises(ValueError):
        _make_multidim_func(dummy_1d_func, n, a, b) # 28.8μs -> 28.9μs (0.498% slower)

def test_one_dim():
    # n is 1d, but a and b are lists of length 1
    n = [4]
    a = [0]
    b = [2]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 90.3μs -> 99.0μs (8.80% slower)
    exp_nodes, exp_weights = dummy_1d_func(4, 0, 2)

# --- Large Scale Test Cases ---

def test_large_3d():
    # n = [10,10,10] (1000 points)
    n = [10,10,10]
    a = [0,0,0]
    b = [1,1,1]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 236μs -> 141μs (66.8% faster)

def test_large_2d_unequal():
    # n = [100, 10] (1000 points)
    n = [100,10]
    a = [0,0]
    b = [1,2]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 168μs -> 105μs (60.0% faster)

def test_large_broadcasted_args():
    # n = [50, 20] (1000), a scalar, b scalar
    n = [50,20]
    a = 0
    b = 1
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 175μs -> 116μs (50.7% faster)

def test_large_high_dim():
    # n = [4,5,5,5] (500 points)
    n = [4,5,5,5]
    a = [0,0,0,0]
    b = [1,1,1,1]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 282μs -> 155μs (81.7% faster)

# --- Additional edge cases ---

def test_nodes_weights_types():
    # Ensure output types are numpy arrays
    n = [2,2]
    a = [0,0]
    b = [1,1]
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 160μs -> 100μs (59.7% faster)

def test_func_return_structure_scalar():
    # If all inputs are scalar, output is tuple of arrays
    n = 2
    a = 0
    b = 1
    nodes, weights = _make_multidim_func(dummy_1d_func, n, a, b) # 59.2μs -> 57.7μs (2.50% faster)

def test_func_return_structure_multi():
    # If inputs are arrays, output is tuple of arrays
    n = [2,2]
    a = [0,0]
    b = [1,1]
    codeflash_output = _make_multidim_func(dummy_1d_func, n, a, b); out = codeflash_output # 158μs -> 100μs (57.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_make_multidim_func-mjvu1q51 and push.