⚡️ Speed up function gridmake by 304%

[codeflash-ai]

📄 304% (3.04x) speedup for gridmake in quantecon/_ce_util.py

⏱️ Runtime : 12.9 milliseconds → 3.18 milliseconds (best of 7 runs)

📝 Explanation and details

The optimized code achieves a 304% speedup by replacing NumPy's high-level array operations (np.tile, np.repeat, np.column_stack) with a Numba JIT-compiled function that uses explicit loops to construct the Cartesian product directly.

Key Optimization

Numba JIT Compilation: The _gridmake2 function is decorated with @njit(cache=True, fastmath=True), which compiles it to native machine code. This eliminates Python interpreter overhead and enables low-level optimizations.

Why This Works:

1. Direct Memory Access: Instead of creating temporary arrays (np.tile creates a full tiled copy, np.repeat creates a full repeated copy), the optimized code writes directly to the output array using explicit loops

2. Memory Efficiency: The original approach allocates multiple intermediate arrays:

   • np.tile(x1, x2.shape[0]) creates a copy of size m*n
   • np.repeat(x2, x1.shape[0]) creates another copy of size m*n
   • np.column_stack creates yet another copy to combine them

   The optimized version allocates just the final output array once

3. Cache Locality: Sequential loop access patterns provide better CPU cache utilization compared to NumPy's strided views and copies

Dtype Promotion Handling: The optimized code correctly handles mixed dtypes (e.g., int + float) by using NumPy's type promotion rules: temp = np.empty(1, dtype=x1.dtype) + np.empty(1, dtype=x2.dtype) determines the result dtype

Performance Characteristics

The speedup is most significant for small-to-medium arrays (10-1000 elements):

• Small arrays: ~170-250% faster (overhead of NumPy operations dominates)
• Large arrays (1000x1000): ~335% faster (JIT compilation benefits compound)
• Multiple vectors (3-5): ~200-320% faster (repeated calls to optimized _gridmake2)

Impact on Workloads

Based on function_references, this function is called in hot paths within quadrature node generation (qnwnorm, _make_multidim_func). These are computationally intensive operations used in numerical integration and computational economics, where:

• gridmake is called inside loops to construct multi-dimensional grids
• Performance scales multiplicatively with problem dimension (d)
• The optimization significantly reduces overhead in multi-dimensional quadrature calculations

The optimization is particularly effective for typical use cases in quantitative economics where grids of 10-100 points per dimension are common.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np
# imports
import pytest
from quantecon._ce_util import gridmake

# unit tests

# ---------- BASIC TEST CASES ----------

def test_two_vectors_simple():
    # Test with two simple 1D arrays
    a = np.array([1, 2])
    b = np.array([3, 4])
    codeflash_output = gridmake(a, b); result = codeflash_output # 59.6μs -> 21.5μs (177% faster)
    expected = np.array([[1, 3],
                         [2, 3],
                         [1, 4],
                         [2, 4]])

def test_three_vectors_simple():
    # Test with three 1D arrays
    a = np.array([1, 2])
    b = np.array([3])
    c = np.array([5, 6])
    codeflash_output = gridmake(a, b, c); result = codeflash_output # 60.3μs -> 19.5μs (210% faster)
    expected = np.array([[1, 3, 5],
                         [2, 3, 5],
                         [1, 3, 6],
                         [2, 3, 6]])

def test_single_element_vectors():
    # Test with vectors of length 1
    a = np.array([7])
    b = np.array([8])
    codeflash_output = gridmake(a, b); result = codeflash_output # 31.7μs -> 12.1μs (162% faster)
    expected = np.array([[7, 8]])

def test_three_single_element_vectors():
    # Test with three vectors of length 1
    a = np.array([1])
    b = np.array([2])
    c = np.array([3])
    codeflash_output = gridmake(a, b, c); result = codeflash_output # 42.7μs -> 16.1μs (165% faster)
    expected = np.array([[1, 2, 3]])

def test_empty_vectors():
    # Test with empty vectors (should produce empty output)
    a = np.array([])
    b = np.array([1, 2])
    codeflash_output = gridmake(a, b); result = codeflash_output # 36.0μs -> 12.1μs (198% faster)
    expected = np.empty((0, 2))

def test_empty_all_vectors():
    # Test with all empty vectors (should produce empty output)
    a = np.array([])
    b = np.array([])
    codeflash_output = gridmake(a, b); result = codeflash_output # 33.7μs -> 12.7μs (165% faster)
    expected = np.empty((0, 2))

def test_different_dtypes():
    # Test with integer and float arrays
    a = np.array([1, 2], dtype=int)
    b = np.array([1.5, 2.5], dtype=float)
    codeflash_output = gridmake(a, b); result = codeflash_output # 38.4μs -> 12.2μs (216% faster)
    expected = np.array([[1, 1.5],
                         [2, 1.5],
                         [1, 2.5],
                         [2, 2.5]])

# ---------- EDGE TEST CASES ----------

def test_high_dimensional_input():
    # Test with 2D input, should raise NotImplementedError
    a = np.array([[1, 2], [3, 4]])
    b = np.array([5, 6])
    with pytest.raises(NotImplementedError):
        gridmake(a, b) # 3.37μs -> 3.85μs (12.5% slower)

def test_mixed_dimensional_input():
    # Test with one 1D and one 2D input, should raise NotImplementedError
    a = np.array([1, 2])
    b = np.array([[3, 4], [5, 6]])
    with pytest.raises(NotImplementedError):
        gridmake(a, b) # 3.40μs -> 3.79μs (10.4% slower)

def test_zero_length_vector():
    # Test with one zero-length vector and one normal vector
    a = np.array([])
    b = np.array([1, 2, 3])
    codeflash_output = gridmake(a, b); result = codeflash_output # 34.8μs -> 12.3μs (184% faster)

def test_three_vectors_one_empty():
    # Test with three vectors, one empty
    a = np.array([1, 2])
    b = np.array([])
    c = np.array([3, 4])
    codeflash_output = gridmake(a, b, c); result = codeflash_output # 50.9μs -> 17.2μs (196% faster)

def test_non_array_input():
    # Test with list instead of np.ndarray
    a = [1, 2]
    b = [3, 4]
    # Should raise AttributeError since lists have no .ndim
    with pytest.raises(AttributeError):
        gridmake(a, b) # 3.33μs -> 3.98μs (16.4% slower)

def test_large_integer_values():
    # Test with large integer values to ensure no overflow
    a = np.array([2**30, 2**31])
    b = np.array([2**32, 2**33])
    codeflash_output = gridmake(a, b); result = codeflash_output # 37.0μs -> 11.7μs (216% faster)
    expected = np.array([[2**30, 2**32],
                         [2**31, 2**32],
                         [2**30, 2**33],
                         [2**31, 2**33]])

def test_negative_values():
    # Test with negative values
    a = np.array([-1, 0, 1])
    b = np.array([-2, 2])
    codeflash_output = gridmake(a, b); result = codeflash_output # 36.3μs -> 11.2μs (224% faster)
    expected = np.array([[-1, -2],
                         [0, -2],
                         [1, -2],
                         [-1, 2],
                         [0, 2],
                         [1, 2]])

def test_non_contiguous_input():
    # Test with non-contiguous input (sliced array)
    a = np.arange(10)[::2]  # [0,2,4,6,8]
    b = np.arange(3)
    codeflash_output = gridmake(a, b); result = codeflash_output # 38.6μs -> 11.7μs (230% faster)
    expected = np.array([[0, 0],
                         [2, 0],
                         [4, 0],
                         [6, 0],
                         [8, 0],
                         [0, 1],
                         [2, 1],
                         [4, 1],
                         [6, 1],
                         [8, 1],
                         [0, 2],
                         [2, 2],
                         [4, 2],
                         [6, 2],
                         [8, 2]])

# ---------- LARGE SCALE TEST CASES ----------

def test_large_vectors():
    # Test with large vectors (1000 elements each)
    a = np.arange(1000)
    b = np.arange(1000)
    codeflash_output = gridmake(a, b); result = codeflash_output # 10.2ms -> 2.34ms (335% faster)

def test_large_three_vectors():
    # Test with three vectors of moderate size (10, 10, 10)
    a = np.arange(10)
    b = np.arange(10)
    c = np.arange(10)
    codeflash_output = gridmake(a, b, c); result = codeflash_output # 76.0μs -> 21.0μs (262% faster)

def test_large_vector_and_single_element():
    # Test with a large vector and a single-element vector
    a = np.arange(1000)
    b = np.array([42])
    codeflash_output = gridmake(a, b); result = codeflash_output # 37.2μs -> 13.7μs (171% faster)
    expected = np.column_stack([a, np.full_like(a, 42)])

def test_large_vector_and_empty():
    # Test with a large vector and an empty vector
    a = np.arange(1000)
    b = np.array([])
    codeflash_output = gridmake(a, b); result = codeflash_output # 40.7μs -> 11.5μs (253% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np
# imports
import pytest  # used for our unit tests
from quantecon._ce_util import gridmake

# unit tests

class TestGridmakeBasic:
    """Basic test cases for fundamental functionality"""
    
    def test_two_simple_vectors(self):
        """Test gridmake with two simple 1D vectors"""
        # Create two simple vectors
        x1 = np.array([1, 2])
        x2 = np.array([3, 4])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 37.5μs -> 10.8μs (247% faster)
    
    def test_three_simple_vectors(self):
        """Test gridmake with three 1D vectors"""
        # Create three simple vectors
        x1 = np.array([1, 2])
        x2 = np.array([3, 4])
        x3 = np.array([5, 6])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3); result = codeflash_output # 54.3μs -> 16.1μs (238% faster)
    
    def test_different_length_vectors(self):
        """Test gridmake with vectors of different lengths"""
        # Create vectors of different lengths
        x1 = np.array([1, 2, 3])
        x2 = np.array([4, 5])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.4μs -> 11.0μs (232% faster)
    
    def test_float_values(self):
        """Test gridmake with floating point values"""
        # Create vectors with float values
        x1 = np.array([1.5, 2.5])
        x2 = np.array([3.5, 4.5])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.3μs -> 11.7μs (211% faster)
    
    def test_negative_values(self):
        """Test gridmake with negative values"""
        # Create vectors with negative values
        x1 = np.array([-1, 0, 1])
        x2 = np.array([-2, 2])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.2μs -> 10.8μs (235% faster)

class TestGridmakeEdgeCases:
    """Edge case tests for unusual conditions"""
    
    def test_single_element_vectors(self):
        """Test gridmake with single-element vectors"""
        # Create single-element vectors
        x1 = np.array([1])
        x2 = np.array([2])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 29.4μs -> 10.9μs (170% faster)
    
    def test_single_element_with_longer_vector(self):
        """Test gridmake with one single-element and one longer vector"""
        # Create vectors of different lengths
        x1 = np.array([1])
        x2 = np.array([2, 3, 4])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.5μs -> 10.8μs (239% faster)
    
    def test_four_vectors(self):
        """Test gridmake with four vectors"""
        # Create four vectors
        x1 = np.array([1, 2])
        x2 = np.array([3, 4])
        x3 = np.array([5, 6])
        x4 = np.array([7, 8])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3, x4); result = codeflash_output # 67.7μs -> 17.4μs (288% faster)
    
    def test_five_vectors(self):
        """Test gridmake with five vectors"""
        # Create five vectors
        x1 = np.array([1, 2])
        x2 = np.array([3, 4])
        x3 = np.array([5, 6])
        x4 = np.array([7, 8])
        x5 = np.array([9, 10])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3, x4, x5); result = codeflash_output # 78.2μs -> 19.2μs (307% faster)
    
    def test_zero_values(self):
        """Test gridmake with zero values"""
        # Create vectors with zeros
        x1 = np.array([0, 0])
        x2 = np.array([0, 1])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 35.8μs -> 10.9μs (230% faster)
    
    def test_very_small_values(self):
        """Test gridmake with very small floating point values"""
        # Create vectors with very small values
        x1 = np.array([1e-10, 2e-10])
        x2 = np.array([3e-10, 4e-10])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.9μs -> 11.3μs (228% faster)
    
    def test_very_large_values(self):
        """Test gridmake with very large values"""
        # Create vectors with large values
        x1 = np.array([1e10, 2e10])
        x2 = np.array([3e10, 4e10])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.5μs -> 11.1μs (228% faster)
    
    def test_mixed_integer_and_float(self):
        """Test gridmake with mixed integer and float arrays"""
        # Create one integer and one float vector
        x1 = np.array([1, 2], dtype=int)
        x2 = np.array([3.5, 4.5], dtype=float)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 37.3μs -> 10.8μs (246% faster)
    
    def test_not_implemented_2d_array(self):
        """Test that 2D arrays raise NotImplementedError"""
        # Create a 2D array
        x1 = np.array([[1, 2], [3, 4]])
        x2 = np.array([5, 6])
        
        # Should raise NotImplementedError
        with pytest.raises(NotImplementedError):
            gridmake(x1, x2) # 3.35μs -> 3.43μs (2.39% slower)
    
    def test_not_implemented_both_2d_arrays(self):
        """Test that two 2D arrays raise NotImplementedError"""
        # Create two 2D arrays
        x1 = np.array([[1, 2], [3, 4]])
        x2 = np.array([[5, 6], [7, 8]])
        
        # Should raise NotImplementedError
        with pytest.raises(NotImplementedError):
            gridmake(x1, x2) # 3.41μs -> 3.43μs (0.758% slower)
    
    def test_asymmetric_three_vectors(self):
        """Test gridmake with three vectors of very different lengths"""
        # Create vectors with different lengths
        x1 = np.array([1])
        x2 = np.array([2, 3])
        x3 = np.array([4, 5, 6, 7])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3); result = codeflash_output # 54.5μs -> 16.0μs (241% faster)
        
        # Verify first column is all 1s
        for i in range(8):
            pass

class TestGridmakeLargeScale:
    """Large scale tests for performance and scalability"""
    
    def test_large_two_vectors(self):
        """Test gridmake with two large vectors"""
        # Create two large vectors (100 elements each)
        x1 = np.arange(100)
        x2 = np.arange(100, 200)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 75.1μs -> 20.2μs (271% faster)
    
    def test_large_three_vectors(self):
        """Test gridmake with three moderately large vectors"""
        # Create three vectors (20 elements each)
        x1 = np.arange(20)
        x2 = np.arange(20, 40)
        x3 = np.arange(40, 60)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3); result = codeflash_output # 133μs -> 63.1μs (111% faster)
    
    def test_many_small_vectors(self):
        """Test gridmake with many small vectors"""
        # Create 10 vectors with 2 elements each
        vectors = [np.array([i, i+1]) for i in range(10)]
        
        # Call gridmake
        codeflash_output = gridmake(*vectors); result = codeflash_output # 165μs -> 39.2μs (321% faster)
        
        # Verify first row
        for i in range(10):
            pass
        
        # Verify last row
        for i in range(10):
            pass
    
    def test_unbalanced_large_vectors(self):
        """Test gridmake with vectors of very different sizes"""
        # Create one small and one large vector
        x1 = np.array([1, 2])
        x2 = np.arange(500)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 39.8μs -> 12.8μs (210% faster)
    
    def test_large_float_vectors(self):
        """Test gridmake with large floating point vectors"""
        # Create two large float vectors
        x1 = np.linspace(0, 1, 100)
        x2 = np.linspace(1, 2, 100)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 56.1μs -> 18.6μs (202% faster)
    
    def test_large_negative_range(self):
        """Test gridmake with large negative value ranges"""
        # Create vectors with large negative ranges
        x1 = np.arange(-500, 0)
        x2 = np.arange(-1000, -900)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 405μs -> 81.7μs (396% faster)
    
    def test_four_medium_vectors(self):
        """Test gridmake with four medium-sized vectors"""
        # Create four vectors (10 elements each)
        x1 = np.arange(10)
        x2 = np.arange(10, 20)
        x3 = np.arange(20, 30)
        x4 = np.arange(30, 40)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3, x4); result = codeflash_output # 171μs -> 84.2μs (104% faster)

class TestGridmakeDataTypes:
    """Test different numpy data types"""
    
    def test_int32_dtype(self):
        """Test gridmake with int32 dtype"""
        # Create vectors with int32 dtype
        x1 = np.array([1, 2], dtype=np.int32)
        x2 = np.array([3, 4], dtype=np.int32)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.3μs -> 12.6μs (188% faster)
    
    def test_int64_dtype(self):
        """Test gridmake with int64 dtype"""
        # Create vectors with int64 dtype
        x1 = np.array([1, 2], dtype=np.int64)
        x2 = np.array([3, 4], dtype=np.int64)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 35.8μs -> 11.1μs (222% faster)
    
    def test_float32_dtype(self):
        """Test gridmake with float32 dtype"""
        # Create vectors with float32 dtype
        x1 = np.array([1.5, 2.5], dtype=np.float32)
        x2 = np.array([3.5, 4.5], dtype=np.float32)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.6μs -> 12.1μs (202% faster)
    
    def test_float64_dtype(self):
        """Test gridmake with float64 dtype"""
        # Create vectors with float64 dtype
        x1 = np.array([1.5, 2.5], dtype=np.float64)
        x2 = np.array([3.5, 4.5], dtype=np.float64)
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.4μs -> 11.0μs (230% faster)

class TestGridmakeCartesianProduct:
    """Test that the cartesian product is computed correctly"""
    
    def test_cartesian_product_order_two_vectors(self):
        """Test the order of cartesian product for two vectors"""
        # Create two vectors
        x1 = np.array([1, 2, 3])
        x2 = np.array([10, 20])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 35.9μs -> 10.9μs (230% faster)
        
        # Verify the complete cartesian product order
        # Expected: (1,10), (2,10), (3,10), (1,20), (2,20), (3,20)
        expected_pairs = [
            (1, 10), (2, 10), (3, 10),
            (1, 20), (2, 20), (3, 20)
        ]
        
        for i, (e1, e2) in enumerate(expected_pairs):
            pass
    
    def test_cartesian_product_order_three_vectors(self):
        """Test the order of cartesian product for three vectors"""
        # Create three small vectors
        x1 = np.array([1, 2])
        x2 = np.array([10, 20])
        x3 = np.array([100, 200])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3); result = codeflash_output # 54.0μs -> 15.3μs (253% faster)
        
        # Verify the complete cartesian product
        # Expected order based on the implementation
        expected = [
            (1, 10, 100), (2, 10, 100),
            (1, 20, 100), (2, 20, 100),
            (1, 10, 200), (2, 10, 200),
            (1, 20, 200), (2, 20, 200)
        ]
        
        for i, (e1, e2, e3) in enumerate(expected):
            pass
    
    def test_all_combinations_present(self):
        """Test that all combinations are present in the output"""
        # Create two vectors
        x1 = np.array([1, 2])
        x2 = np.array([3, 4, 5])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 36.3μs -> 10.8μs (236% faster)
        
        # Create set of all combinations
        result_set = set()
        for i in range(result.shape[0]):
            result_set.add((result[i, 0], result[i, 1]))
        
        # Expected combinations
        expected_set = {(1, 3), (1, 4), (1, 5), (2, 3), (2, 4), (2, 5)}
        for combo in expected_set:
            pass
    
    def test_no_duplicate_rows(self):
        """Test that there are no duplicate rows in the output"""
        # Create two vectors
        x1 = np.array([1, 2, 3])
        x2 = np.array([4, 5, 6])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2); result = codeflash_output # 35.9μs -> 10.9μs (229% faster)
        
        # Convert rows to tuples and check for duplicates
        rows_as_tuples = [tuple(result[i, :]) for i in range(result.shape[0])]
    
    def test_row_count_matches_product(self):
        """Test that the number of rows equals the product of input lengths"""
        # Create vectors of various lengths
        x1 = np.array([1, 2, 3, 4, 5])
        x2 = np.array([6, 7, 8])
        x3 = np.array([9, 10])
        
        # Call gridmake
        codeflash_output = gridmake(x1, x2, x3); result = codeflash_output # 53.5μs -> 15.8μs (238% faster)
        
        # Verify row count equals product of lengths
        expected_rows = len(x1) * len(x2) * len(x3)
# 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-gridmake-mjvywofc and push.