⚡️ Speed up function mlinspace by 42%

[codeflash-ai]

📄 42% (0.42x) speedup for mlinspace in quantecon/_gridtools.py

⏱️ Runtime : 1.95 milliseconds → 1.37 milliseconds (best of 117 runs)

📝 Explanation and details

The optimized code achieves a 42% speedup through several key optimizations that reduce overhead and improve memory efficiency:

1. Early exit optimizations for edge cases:

• Added checks for n == 0 (empty input) and n == 1 (single dimension) cases that bypass expensive computation and directly return results. This is particularly effective for 1D grids, showing 60-87% speedups in test cases.

2. Replaced np.prod() with manual multiplication:

• Changed from l = np.prod(shapes) to a simple loop for dim in shapes: l *= dim. This avoids creating intermediate arrays and function call overhead for a scalar result.

3. Optimized repetitions calculation:

• Eliminated list operations ([1] + shapes[:-1], .reverse(), .tolist()) and replaced with direct NumPy array allocation and in-place computation using accumulators. This removes unnecessary memory allocations and copying.

4. Memory allocation improvements:

• Changed from np.zeros() to np.empty() for the output array since values will be overwritten anyway, saving initialization time.
• Pre-allocated repetitions as np.int64 arrays instead of using Python lists.

5. Minor enhancements in mlinspace:

• Added explicit order='C' parameter to np.asarray calls for better memory layout consistency.
• Used nums.shape[0] instead of len(nums) for slight efficiency gain.

The optimizations are most effective for:

• 1D cases (60-87% faster): Early exit path avoids all cartesian product computation
• Small to medium grids (30-50% faster): Overhead reductions are more significant relative to total runtime
• All dimensionalities: The repetitions calculation improvements benefit both C and F order layouts consistently

These changes maintain identical behavior while eliminating computational bottlenecks in the setup phase before the core _repeat_1d loop (which remains 99%+ of total runtime).

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_gridtools.py::test_mlinsplace	54.3μs	38.3μs	41.8%✅

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from numba import njit
from quantecon._gridtools import mlinspace

# unit tests

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

def test_1d_basic():
    # 1D grid, 5 points from 0 to 1
    codeflash_output = mlinspace([0], [1], [5]); result = codeflash_output # 42.0μs -> 23.6μs (78.2% faster)
    expected = np.linspace(0, 1, 5).reshape(-1, 1)

def test_2d_basic():
    # 2D grid, 3 points from 0 to 1 in x, 2 points from 10 to 20 in y
    codeflash_output = mlinspace([0, 10], [1, 20], [3, 2]); result = codeflash_output # 48.5μs -> 34.9μs (38.8% faster)
    expected_x = np.linspace(0, 1, 3)
    expected_y = np.linspace(10, 20, 2)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])

def test_3d_basic():
    # 3D grid, 2 points in each dimension
    codeflash_output = mlinspace([0, 0, 0], [1, 2, 3], [2, 2, 2]); result = codeflash_output # 54.5μs -> 38.8μs (40.7% faster)
    expected_x = np.linspace(0, 1, 2)
    expected_y = np.linspace(0, 2, 2)
    expected_z = np.linspace(0, 3, 2)
    expected = np.array([[x, y, z] for x in expected_x for y in expected_y for z in expected_z])

def test_order_F_basic():
    # 2D grid, Fortran order
    codeflash_output = mlinspace([0, 10], [1, 20], [3, 2], order='F'); result = codeflash_output # 51.1μs -> 35.0μs (45.9% faster)
    expected_x = np.linspace(0, 1, 3)
    expected_y = np.linspace(10, 20, 2)
    expected = np.array([[x, y] for y in expected_y for x in expected_x])

def test_non_integer_bounds():
    # Non-integer bounds
    codeflash_output = mlinspace([0.5, 2.5], [1.5, 3.5], [2, 3]); result = codeflash_output # 49.4μs -> 33.9μs (45.7% faster)
    expected_x = np.linspace(0.5, 1.5, 2)
    expected_y = np.linspace(2.5, 3.5, 3)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])

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

def test_single_point_per_dim():
    # Each dimension has only one point
    codeflash_output = mlinspace([1, 2, 3], [1, 2, 3], [1, 1, 1]); result = codeflash_output # 57.3μs -> 43.5μs (31.7% faster)
    expected = np.array([[1, 2, 3]])


def test_negative_num_nodes():
    # Negative number of nodes (should raise ValueError)
    with pytest.raises(ValueError):
        mlinspace([0], [1], [-5]) # 17.0μs -> 17.7μs (3.79% slower)



def test_large_numbers():
    # Large numbers for bounds
    codeflash_output = mlinspace([1e10, -1e10], [1e10+1, -1e10+1], [2, 2]); result = codeflash_output # 80.3μs -> 60.3μs (33.1% faster)
    expected_x = np.linspace(1e10, 1e10+1, 2)
    expected_y = np.linspace(-1e10, -1e10+1, 2)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])

def test_reverse_bounds():
    # Lower bound > upper bound
    codeflash_output = mlinspace([1, 2], [0, 1], [2, 2]); result = codeflash_output # 55.6μs -> 37.5μs (48.1% faster)
    expected_x = np.linspace(1, 0, 2)
    expected_y = np.linspace(2, 1, 2)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])


def test_non_array_inputs():
    # Inputs are lists, not arrays
    codeflash_output = mlinspace([0, 1], [1, 2], [2, 2]); result = codeflash_output # 51.5μs -> 35.2μs (46.2% faster)
    expected_x = np.linspace(0, 1, 2)
    expected_y = np.linspace(1, 2, 2)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])


def test_large_1d():
    # Large 1D grid, 1000 points
    codeflash_output = mlinspace([0], [1], [1000]); result = codeflash_output # 75.3μs -> 46.9μs (60.7% faster)
    expected = np.linspace(0, 1, 1000).reshape(-1, 1)

def test_large_2d():
    # Large 2D grid, 100 x 10 points
    codeflash_output = mlinspace([0, 0], [1, 1], [100, 10]); result = codeflash_output # 62.1μs -> 45.8μs (35.4% faster)
    expected_x = np.linspace(0, 1, 100)
    expected_y = np.linspace(0, 1, 10)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])

def test_large_3d():
    # Large 3D grid, 10 x 10 x 10 points
    codeflash_output = mlinspace([0, 0, 0], [1, 1, 1], [10, 10, 10]); result = codeflash_output # 71.1μs -> 52.8μs (34.7% faster)
    expected_x = np.linspace(0, 1, 10)
    expected_y = np.linspace(0, 1, 10)
    expected_z = np.linspace(0, 1, 10)
    expected = np.array([[x, y, z] for x in expected_x for y in expected_y for z in expected_z])

def test_large_order_F():
    # Large 2D grid, Fortran order
    codeflash_output = mlinspace([0, 0], [1, 1], [100, 10], order='F'); result = codeflash_output # 62.0μs -> 47.0μs (32.0% faster)
    expected_x = np.linspace(0, 1, 100)
    expected_y = np.linspace(0, 1, 10)
    expected = np.array([[x, y] for y in expected_y for x in expected_x])

def test_large_float_bounds():
    # Large grid with float bounds
    codeflash_output = mlinspace([0.1, 2.2], [9.9, 3.3], [100, 10]); result = codeflash_output # 63.7μs -> 43.0μs (48.2% faster)
    expected_x = np.linspace(0.1, 9.9, 100)
    expected_y = np.linspace(2.2, 3.3, 10)
    expected = np.array([[x, y] for x in expected_x for y in expected_y])
# 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 numba import njit
from quantecon._gridtools import mlinspace

# unit tests

# ---------------------------
# Basic Test Cases
# ---------------------------

def test_single_dimension_basic():
    # 1D: a=0, b=1, nums=5
    codeflash_output = mlinspace([0], [1], [5]); result = codeflash_output # 60.4μs -> 36.3μs (66.5% faster)
    expected = np.linspace(0, 1, 5).reshape(-1, 1)

def test_two_dimensions_basic():
    # 2D: a=[0,0], b=[1,1], nums=[2,3]
    codeflash_output = mlinspace([0,0], [1,1], [2,3]); result = codeflash_output # 55.2μs -> 39.6μs (39.5% faster)
    # Should produce 2x3 = 6 points
    expected = np.array([[0,0],[0,0.5],[0,1],[1,0],[1,0.5],[1,1]])

def test_three_dimensions_basic():
    # 3D: a=[0,0,0], b=[1,2,3], nums=[2,2,2]
    codeflash_output = mlinspace([0,0,0], [1,2,3], [2,2,2]); result = codeflash_output # 56.9μs -> 41.0μs (38.8% faster)

def test_order_F_vs_C():
    # Test 'F' order
    codeflash_output = mlinspace([0,0], [1,1], [2,2], order='C'); result_C = codeflash_output # 51.4μs -> 35.7μs (44.1% faster)
    codeflash_output = mlinspace([0,0], [1,1], [2,2], order='F'); result_F = codeflash_output # 27.7μs -> 17.5μs (58.1% faster)

def test_non_integer_bounds():
    # Non-integer bounds
    codeflash_output = mlinspace([0.1, 2.5], [1.1, 3.5], [2, 2]); result = codeflash_output # 47.7μs -> 33.5μs (42.3% faster)
    expected = np.array([[0.1,2.5],[0.1,3.5],[1.1,2.5],[1.1,3.5]])

# ---------------------------
# Edge Test Cases
# ---------------------------


def test_one_point_per_dimension():
    # nums = [1,1,...]: should produce a single point
    codeflash_output = mlinspace([2,3,4], [2,3,4], [1,1,1]); result = codeflash_output # 83.9μs -> 65.2μs (28.8% faster)
    expected = np.array([[2,3,4]])

def test_negative_nums():
    # nums contains negative: should raise ValueError from np.linspace
    with pytest.raises(ValueError):
        mlinspace([0], [1], [-1]) # 10.5μs -> 10.5μs (0.219% slower)
    with pytest.raises(ValueError):
        mlinspace([0,0], [1,1], [2,-5]) # 22.6μs -> 23.3μs (2.81% slower)


def test_non_numeric_input():
    # Non-numeric input: should raise TypeError from np.asarray or np.linspace
    with pytest.raises(ValueError):
        mlinspace(['a'], ['b'], [2]) # 9.66μs -> 9.72μs (0.638% slower)
    with pytest.raises(ValueError):
        mlinspace([0], [1], ['x']) # 6.02μs -> 5.84μs (3.10% faster)

def test_large_range():
    # Very large range: check for correct output, no overflow
    codeflash_output = mlinspace([1e10], [1e10+1], [2]); result = codeflash_output # 67.4μs -> 40.3μs (67.2% faster)
    expected = np.array([[1e10],[1e10+1]])

def test_reverse_bounds():
    # a > b: np.linspace allows this, should produce descending output
    codeflash_output = mlinspace([1], [0], [3]); result = codeflash_output # 50.7μs -> 28.0μs (81.0% faster)
    expected = np.array([[1],[0.5],[0]])


def test_large_negative_bounds():
    # Large negative bounds
    codeflash_output = mlinspace([-1e5, -1e6], [-1e4, -1e3], [2,2]); result = codeflash_output # 80.0μs -> 60.2μs (33.0% faster)
    expected = np.array([[-1e5, -1e6], [-1e5, -1e3], [-1e4, -1e6], [-1e4, -1e3]])

# ---------------------------
# Large Scale Test Cases
# ---------------------------

def test_large_1d():
    # 1D, nums=1000
    codeflash_output = mlinspace([0], [10], [1000]); result = codeflash_output # 53.7μs -> 28.7μs (87.4% faster)

def test_large_2d():
    # 2D, nums=[30,30] (900 points)
    codeflash_output = mlinspace([0,0], [1,1], [30,30]); result = codeflash_output # 58.1μs -> 42.4μs (37.2% faster)

def test_large_3d():
    # 3D, nums=[10,10,10] (1000 points)
    codeflash_output = mlinspace([0,0,0], [1,1,1], [10,10,10]); result = codeflash_output # 61.9μs -> 46.9μs (32.1% faster)

def test_large_values():
    # Large values in a and b, but small nums
    codeflash_output = mlinspace([1e8,1e9], [1e8+1,1e9+1], [2,2]); result = codeflash_output # 51.1μs -> 35.6μs (43.7% faster)
    expected = np.array([[1e8,1e9],[1e8,1e9+1],[1e8+1,1e9],[1e8+1,1e9+1]])

def test_large_scale_order_F():
    # Large scale with order 'F'
    codeflash_output = mlinspace([0,0,0], [1,1,1], [10,10,10], order='F'); result = codeflash_output # 63.4μs -> 45.9μs (38.1% faster)
    # Should contain same points as order 'C'
    codeflash_output = mlinspace([0,0,0], [1,1,1], [10,10,10], order='C'); result_C = codeflash_output # 36.7μs -> 26.0μs (41.5% faster)

def test_large_scale_non_uniform_nums():
    # Large scale with non-uniform nums
    codeflash_output = mlinspace([0,0,0], [1,2,3], [5,10,2]); result = codeflash_output # 56.2μs -> 40.8μs (37.7% 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-mlinspace-mggvnurl and push.