⚡️ Speed up function _gridmake2 by 106%

[codeflash-ai]

📄 106% (1.06x) speedup for _gridmake2 in quantecon/_ce_util.py

⏱️ Runtime : 2.33 milliseconds → 1.13 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 105% speedup by replacing NumPy's high-level array operations (np.tile, np.repeat, np.column_stack) with Numba JIT-compiled loops that directly construct the output array.

Key Optimizations

1. Numba JIT Compilation (@njit)
Two specialized helper functions (_gridmake2_1d_1d and _gridmake2_nd_1d) are decorated with @njit, enabling machine-code compilation. This eliminates Python interpreter overhead and enables CPU-level optimizations.

2. Direct Memory Allocation
Instead of creating intermediate arrays with np.tile and np.repeat, then combining them with np.column_stack, the optimized code:

• Pre-allocates the exact output array size once
• Fills it directly via slicing operations in a loop
• Avoids multiple temporary array allocations and copies

3. Efficient Memory Access Pattern
The loop iterates over x2 elements, writing contiguous blocks of x1 values. This provides good cache locality since sequential memory writes are efficient.

Performance Characteristics

Test Results Show:

• Small arrays (2-3 elements): 85-150% faster - JIT compilation overhead is amortized by eliminating temporary arrays
• Larger arrays (100-1000 elements): 100-115% faster - Cache-friendly sequential writes scale well
• Edge cases (empty arrays, single elements): 70-150% faster - Pre-allocation still beats NumPy's internal logic

Why This Works:

• np.column_stack internally copies data multiple times to build the output
• np.tile and np.repeat each allocate temporary arrays
• The optimized version does one allocation and direct assignment, reducing memory bandwidth pressure

Impact on Workloads

From function_references, _gridmake2 is called by gridmake(), which creates cartesian products for computational economics applications. The function appears to be a building block for grid construction, likely used in:

• Dynamic programming state space enumeration
• Monte Carlo simulations parameter combinations
• Numerical optimization over multi-dimensional grids

Since gridmake may call _gridmake2 multiple times for >2 arrays (chained calls in the loop), the 105% speedup compounds when constructing large grids, making this optimization particularly valuable for high-dimensional problems common in quantitative economics.

The optimization is most beneficial when _gridmake2 is called repeatedly with medium-to-large arrays, which is typical in grid-based computational methods.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numba
import numpy as np
# imports
import pytest
from quantecon._ce_util import _gridmake2

# unit tests

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

def test_basic_1d_vectors():
    # Test with simple 1D vectors
    x1 = np.array([1, 2])
    x2 = np.array([3, 4])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 57.7μs -> 28.6μs (102% faster)
    # Expected: cartesian product, each row is [x1_i, x2_j]
    expected = np.array([
        [1, 3],
        [2, 3],
        [1, 4],
        [2, 4]
    ])

def test_basic_1d_vectors_different_lengths():
    # x1 longer than x2
    x1 = np.array([1, 2, 3])
    x2 = np.array([4, 5])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 42.7μs -> 19.5μs (119% faster)
    expected = np.array([
        [1, 4],
        [2, 4],
        [3, 4],
        [1, 5],
        [2, 5],
        [3, 5]
    ])

def test_basic_1d_vectors_single_element():
    # One element in each vector
    x1 = np.array([7])
    x2 = np.array([8])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 29.5μs -> 15.8μs (86.2% faster)
    expected = np.array([[7, 8]])

def test_basic_2d_x1_1d_x2():
    # x1 is a 2D array, x2 is 1D
    x1 = np.array([[1, 2], [3, 4]])
    x2 = np.array([5, 6])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 35.1μs -> 16.7μs (110% faster)
    # Each row of x1 is repeated for each element of x2,
    # and x2 is repeated for each row of x1
    expected = np.array([
        [1, 2, 5],
        [3, 4, 5],
        [1, 2, 6],
        [3, 4, 6]
    ])

def test_basic_2d_x1_1d_x2_single_row():
    # x1 is 2D with one row, x2 is 1D
    x1 = np.array([[10, 20]])
    x2 = np.array([30, 40])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.2μs -> 15.1μs (127% faster)
    expected = np.array([
        [10, 20, 30],
        [10, 20, 40]
    ])

def test_basic_2d_x1_1d_x2_single_col():
    # x1 is 2D with two rows, one column
    x1 = np.array([[1], [2]])
    x2 = np.array([3, 4])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.9μs -> 14.3μs (138% faster)
    expected = np.array([
        [1, 3],
        [2, 3],
        [1, 4],
        [2, 4]
    ])

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

def test_empty_x1_1d():
    # x1 is empty 1D, x2 is 1D
    x1 = np.array([])
    x2 = np.array([1, 2])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.6μs -> 17.0μs (103% faster)

def test_empty_x2_1d():
    # x2 is empty 1D, x1 is 1D
    x1 = np.array([1, 2])
    x2 = np.array([])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 36.3μs -> 15.6μs (132% faster)

def test_empty_x1_2d():
    # x1 is empty 2D, x2 is 1D
    x1 = np.empty((0, 2))
    x2 = np.array([1, 2])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 30.5μs -> 15.8μs (93.5% faster)

def test_empty_x2_1d_with_2d_x1():
    # x2 is empty 1D, x1 is 2D
    x1 = np.array([[1, 2], [3, 4]])
    x2 = np.array([])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.1μs -> 15.8μs (117% faster)

def test_single_element_2d_x1():
    # x1 is 2D, single row, x2 is 1D, single element
    x1 = np.array([[1, 2]])
    x2 = np.array([3])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 26.5μs -> 14.2μs (86.5% faster)
    expected = np.array([[1, 2, 3]])

def test_single_element_2d_x1_multiple_cols():
    # x1 is 2D, single row, multiple columns, x2 is 1D, single element
    x1 = np.array([[9, 8, 7]])
    x2 = np.array([6])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 26.3μs -> 14.4μs (82.7% faster)
    expected = np.array([[9, 8, 7, 6]])

def test_not_implemented_error():
    # x2 is 2D (should raise NotImplementedError)
    x1 = np.array([1, 2])
    x2 = np.array([[3, 4], [5, 6]])
    with pytest.raises(NotImplementedError):
        _gridmake2(x1, x2) # 2.19μs -> 6.91μs (68.3% slower)

def test_not_implemented_error_both_2d():
    # Both x1 and x2 are 2D (should raise NotImplementedError)
    x1 = np.array([[1, 2], [3, 4]])
    x2 = np.array([[5, 6], [7, 8]])
    with pytest.raises(NotImplementedError):
        _gridmake2(x1, x2) # 2.38μs -> 7.28μs (67.4% slower)

def test_dtype_preservation():
    # Check that output dtype matches input dtype (float)
    x1 = np.array([1.1, 2.2])
    x2 = np.array([3.3, 4.4])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 36.3μs -> 15.3μs (137% faster)

def test_dtype_preservation_int():
    # Check that output dtype matches input dtype (int)
    x1 = np.array([1, 2], dtype=np.int32)
    x2 = np.array([3, 4], dtype=np.int32)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.9μs -> 16.1μs (117% faster)

def test_large_numbers():
    # Use large numbers to check for overflow
    x1 = np.array([2**30, 2**31], dtype=np.int64)
    x2 = np.array([2**32, 2**33], dtype=np.int64)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.3μs -> 14.1μs (142% faster)
    expected = np.array([
        [2**30, 2**32],
        [2**31, 2**32],
        [2**30, 2**33],
        [2**31, 2**33]
    ], dtype=np.int64)

def test_negative_numbers():
    # Test with negative numbers
    x1 = np.array([-1, -2])
    x2 = np.array([-3, -4])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.9μs -> 13.9μs (152% faster)
    expected = np.array([
        [-1, -3],
        [-2, -3],
        [-1, -4],
        [-2, -4]
    ])

def test_mixed_sign_numbers():
    # Test with mix of positive and negative
    x1 = np.array([-1, 2])
    x2 = np.array([3, -4])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.8μs -> 14.0μs (148% faster)
    expected = np.array([
        [-1, 3],
        [2, 3],
        [-1, -4],
        [2, -4]
    ])

def test_float_precision():
    # Test with floats that could lose precision
    x1 = np.array([1.123456789, 2.987654321])
    x2 = np.array([3.555555555, 4.444444444])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 35.1μs -> 14.0μs (151% faster)
    expected = np.array([
        [1.123456789, 3.555555555],
        [2.987654321, 3.555555555],
        [1.123456789, 4.444444444],
        [2.987654321, 4.444444444]
    ])

def test_x1_2d_x2_1d_zero_row_x1():
    # x1 is 2D, zero rows, x2 is 1D
    x1 = np.empty((0, 3))
    x2 = np.array([1, 2])
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 30.7μs -> 14.8μs (108% faster)

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

def test_large_scale_1d():
    # Large 1D arrays
    x1 = np.arange(100)
    x2 = np.arange(50)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 49.1μs -> 22.9μs (115% faster)

def test_large_scale_2d_x1_1d_x2():
    # Large 2D x1, 1D x2
    x1 = np.arange(100).reshape(50, 2)
    x2 = np.arange(20)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 43.9μs -> 20.8μs (111% faster)

def test_large_scale_2d_x1_1d_x2_single_col():
    # Large 2D x1 (100 rows, 1 col), x2 1D (10 elements)
    x1 = np.arange(100).reshape(100, 1)
    x2 = np.arange(10)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 36.7μs -> 19.9μs (84.2% faster)

def test_large_scale_empty_x1():
    # Large x2, empty x1
    x1 = np.empty((0,))
    x2 = np.arange(500)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.8μs -> 19.8μs (70.8% faster)

def test_large_scale_empty_x2():
    # Large x1, empty x2
    x1 = np.arange(500)
    x2 = np.empty((0,))
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 36.2μs -> 14.5μs (150% faster)

def test_large_scale_2d_x1_1d_x2_empty_x1():
    # Large 2D x1, empty x1, x2 1D
    x1 = np.empty((0, 10))
    x2 = np.arange(100)
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 31.3μs -> 16.2μs (93.6% faster)

def test_large_scale_2d_x1_1d_x2_empty_x2():
    # Large 2D x1, x2 empty
    x1 = np.arange(1000).reshape(100, 10)
    x2 = np.empty((0,))
    codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.2μs -> 14.8μs (131% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

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

# unit tests

class TestBasic1DVectors:
    """Test basic functionality with two 1D vectors"""
    
    def test_simple_equal_length_vectors(self):
        # Test with two simple 1D arrays of equal length
        x1 = np.array([1, 2, 3])
        x2 = np.array([4, 5, 6])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 35.2μs -> 14.7μs (140% faster)
        
        # Check that first column contains x1 tiled 3 times
        expected_col1 = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3])
        
        # Check that second column contains x2 repeated 3 times each
        expected_col2 = np.array([4, 4, 4, 5, 5, 5, 6, 6, 6])
    
    def test_different_length_vectors(self):
        # Test with two 1D arrays of different lengths
        x1 = np.array([1, 2])
        x2 = np.array([10, 20, 30])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.6μs -> 14.3μs (143% faster)
        
        # First column: x1 tiled 3 times
        expected_col1 = np.array([1, 2, 1, 2, 1, 2])
        
        # Second column: x2 repeated 2 times each
        expected_col2 = np.array([10, 10, 20, 20, 30, 30])
    
    def test_single_element_arrays(self):
        # Test with single element arrays
        x1 = np.array([5])
        x2 = np.array([10])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 28.0μs -> 14.3μs (96.0% faster)
    
    def test_single_element_first_multiple_second(self):
        # Test with single element in x1, multiple in x2
        x1 = np.array([7])
        x2 = np.array([1, 2, 3, 4])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.9μs -> 14.2μs (145% faster)
    
    def test_negative_numbers(self):
        # Test with negative numbers
        x1 = np.array([-1, -2])
        x2 = np.array([-10, -20])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.4μs -> 14.1μs (144% faster)
        
        # Check values
        expected_col1 = np.array([-1, -2, -1, -2])
        expected_col2 = np.array([-10, -10, -20, -20])
    
    def test_floating_point_numbers(self):
        # Test with floating point numbers
        x1 = np.array([1.5, 2.5])
        x2 = np.array([0.1, 0.2, 0.3])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 35.3μs -> 15.8μs (124% faster)
        
        # Check values with floating point comparison
        expected_col1 = np.array([1.5, 2.5, 1.5, 2.5, 1.5, 2.5])
        expected_col2 = np.array([0.1, 0.1, 0.2, 0.2, 0.3, 0.3])
    
    def test_zeros_in_arrays(self):
        # Test with zeros in the arrays
        x1 = np.array([0, 1, 2])
        x2 = np.array([0, 0])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.7μs -> 14.3μs (143% faster)
        
        # Check that zeros are handled correctly
        expected_col1 = np.array([0, 1, 2, 0, 1, 2])
        expected_col2 = np.array([0, 0, 0, 0, 0, 0])
    
    def test_duplicate_values(self):
        # Test with duplicate values in arrays
        x1 = np.array([1, 1, 2])
        x2 = np.array([3, 3])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.4μs -> 14.4μs (140% faster)
        
        # Check values
        expected_col1 = np.array([1, 1, 2, 1, 1, 2])
        expected_col2 = np.array([3, 3, 3, 3, 3, 3])

class TestMatrix2DWith1DVector:
    """Test functionality with 2D matrix as x1 and 1D vector as x2"""
    
    def test_2d_matrix_with_1d_vector(self):
        # Test with x1 as 2D matrix and x2 as 1D vector
        x1 = np.array([[1, 2], [3, 4]])  # 2x2 matrix
        x2 = np.array([10, 20])  # 1D vector of length 2
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.4μs -> 14.5μs (130% faster)
        
        # First two columns should be x1 tiled
        # Expected: [[1,2], [3,4], [1,2], [3,4]]
        expected_first_cols = np.array([[1, 2], [3, 4], [1, 2], [3, 4]])
        
        # Last column should be x2 repeated for each row of x1
        expected_last_col = np.array([10, 10, 20, 20])
    
    def test_2d_matrix_single_column(self):
        # Test with x1 as 2D matrix with single column
        x1 = np.array([[1], [2], [3]])  # 3x1 matrix
        x2 = np.array([100, 200])  # 1D vector
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.4μs -> 14.3μs (134% faster)
        
        # First column should be x1 tiled
        expected_col1 = np.array([1, 2, 3, 1, 2, 3])
        
        # Second column should be x2 repeated
        expected_col2 = np.array([100, 100, 100, 200, 200, 200])
    
    def test_2d_matrix_single_row(self):
        # Test with x1 as 2D matrix with single row
        x1 = np.array([[1, 2, 3]])  # 1x3 matrix
        x2 = np.array([10, 20])  # 1D vector
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.2μs -> 14.2μs (134% faster)
        
        # First three columns should be x1 tiled
        expected_first_cols = np.array([[1, 2, 3], [1, 2, 3]])
        
        # Last column should be x2 repeated
        expected_last_col = np.array([10, 20])
    
    def test_2d_matrix_multiple_columns_multiple_rows(self):
        # Test with larger 2D matrix
        x1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # 3x3 matrix
        x2 = np.array([100])  # Single element vector
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 26.3μs -> 14.3μs (84.5% faster)
        
        # Last column should be all 100s
        expected_last_col = np.array([100, 100, 100])
    
    def test_2d_matrix_with_floats(self):
        # Test with floating point 2D matrix
        x1 = np.array([[1.1, 2.2], [3.3, 4.4]])
        x2 = np.array([0.5, 1.5])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.5μs -> 16.0μs (109% faster)
        
        # Check with floating point tolerance
        expected_first_cols = np.array([[1.1, 2.2], [3.3, 4.4], [1.1, 2.2], [3.3, 4.4]])
        
        expected_last_col = np.array([0.5, 0.5, 1.5, 1.5])

class TestEdgeCases:
    """Test edge cases and boundary conditions"""
    
    def test_very_small_values(self):
        # Test with very small floating point values
        x1 = np.array([1e-10, 2e-10])
        x2 = np.array([1e-15, 2e-15])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.7μs -> 15.0μs (132% faster)
        
        # Check that small values are preserved
        expected_col1 = np.array([1e-10, 2e-10, 1e-10, 2e-10])
        expected_col2 = np.array([1e-15, 1e-15, 2e-15, 2e-15])
    
    def test_large_values(self):
        # Test with large values
        x1 = np.array([1e10, 2e10])
        x2 = np.array([1e15])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 27.5μs -> 14.7μs (87.7% faster)
        
        # Check values
        expected_col1 = np.array([1e10, 2e10])
        expected_col2 = np.array([1e15, 1e15])
    
    def test_mixed_positive_negative_zero(self):
        # Test with mix of positive, negative, and zero
        x1 = np.array([-5, 0, 5])
        x2 = np.array([-1, 0, 1])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 34.3μs -> 14.4μs (138% faster)
        
        # Check values
        expected_col1 = np.array([-5, 0, 5, -5, 0, 5, -5, 0, 5])
        expected_col2 = np.array([-1, -1, -1, 0, 0, 0, 1, 1, 1])
    
    def test_integer_dtype_preservation(self):
        # Test that integer dtypes are handled correctly
        x1 = np.array([1, 2, 3], dtype=np.int32)
        x2 = np.array([10, 20], dtype=np.int64)
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 35.6μs -> 15.6μs (129% faster)
        
        # Check values
        expected_col1 = np.array([1, 2, 3, 1, 2, 3])
        expected_col2 = np.array([10, 10, 10, 20, 20, 20])
    
    def test_2d_matrix_single_element(self):
        # Test with 2D matrix containing single element
        x1 = np.array([[5]])  # 1x1 matrix
        x2 = np.array([10, 20, 30])
        
        # Call the function
        codeflash_output = _gridmake2(x1, x2); result = codeflash_output # 33.2μs -> 14.4μs (130% faster)

class TestErrorCases:
    """Test cases that should raise errors"""
    
    def test_both_2d_raises_error(self):
        # Test that both inputs being 2D raises NotImplementedError
        x1 = np.array([[1, 2], [3, 4]])
        x2 = np.array([[5, 6], [7, 8]])
        
        # Should raise NotImplementedError
        with pytest.raises(NotImplementedError):
            _gridmake2(x1, x2) # 2.38μs -> 6.84μs (65.3% slower)
    
    def test_x1_1d_x2_2d_raises_error(self):
        # Test that x1 as 1D and x2 as 2D raises NotImplementedError
        x1 = np.array([1, 2, 3])
        x2 = np.array([[4, 5], [6, 7]])
        
        # Should raise NotImplementedError
        with pytest.raises(NotImplementedError):
            _gridmake2(x1, x2) # 2.40μs -> 7.10μs (66.3% slower)
    
    

To edit these changes git checkout codeflash/optimize-_gridmake2-mjvz595g and push.