⚡️ Speed up function linear_palette by 281%

[codeflash-ai]

📄 281% (2.81x) speedup for linear_palette in src/bokeh/palettes.py

⏱️ Runtime : 1.19 milliseconds → 312 microseconds (best of 195 runs)

📝 Explanation and details

The optimization achieves a 280% speedup by eliminating the expensive NumPy dependency and replacing it with pure Python math operations.

Key optimizations:

1. Removed NumPy import and np.linspace() call: The original code used np.linspace(0, len(palette)-1, num=n) which involves NumPy array creation and floating-point operations. The optimized version replaces this with simple arithmetic: step = (len(palette) - 1) / (n - 1) and range(n).

2. Added early return for n=1: When only one color is requested, the function now returns immediately with (palette[0],) instead of going through the full computation pipeline.

3. Simplified index calculation: Instead of generating a NumPy array of evenly-spaced values, the code now computes indices on-demand using step * i for each iteration.

Why this is faster:

• NumPy overhead elimination: NumPy imports are expensive and np.linspace() creates intermediate arrays that consume memory and CPU cycles
• Reduced function call overhead: range(n) is a built-in Python iterator that's much faster than NumPy array operations for small to medium-sized palettes
• Branch optimization: The n==1 special case avoids all mathematical computation for a common use case

Performance characteristics:
The optimization shows dramatic improvements for small values of n (300-5000% faster for n=1 cases) and consistent 60-900% improvements across all test cases. It's particularly effective for typical palette operations where n is small to moderate, which represents the majority of real-world usage patterns in data visualization libraries.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Tuple

# imports
import pytest  # used for our unit tests
from bokeh.palettes import linear_palette

# unit tests

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

def test_basic_evenly_spaced_palette():
    # Test with a palette of 5 colors, requesting 3 colors
    palette = ("#000000", "#111111", "#222222", "#333333", "#444444")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 30.4μs -> 2.98μs (921% faster)

def test_basic_n_equals_palette_length():
    # Requesting n equal to palette length should return the palette itself
    palette = ("#ff0000", "#00ff00", "#0000ff")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 22.1μs -> 2.43μs (808% faster)

def test_basic_n_is_1():
    # Requesting 1 color should return the first color
    palette = ("#abcabc", "#defdef", "#123123")
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 23.4μs -> 674ns (3369% faster)

def test_basic_n_is_2():
    # Requesting 2 colors should return first and last
    palette = ("#a", "#b", "#c", "#d")
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 21.1μs -> 2.52μs (739% faster)

def test_basic_non_hex_strings():
    # Palette can contain any strings, not just hex codes
    palette = ("red", "green", "blue", "yellow")
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 20.0μs -> 2.30μs (770% faster)

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

def test_edge_n_zero():
    # Requesting 0 colors should return empty tuple
    palette = ("#000", "#111", "#222")
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 17.1μs -> 1.93μs (789% faster)

def test_edge_palette_empty():
    # Empty palette should always return empty tuple
    palette = ()
    codeflash_output = linear_palette(palette, 0); result = codeflash_output
    # Requesting n > 0 should not raise, just return empty
    codeflash_output = linear_palette(palette, 1); result = codeflash_output

def test_edge_n_greater_than_palette_length():
    # Should raise ValueError if n > len(palette)
    palette = ("#a", "#b")
    with pytest.raises(ValueError):
        linear_palette(palette, 3) # 1.49μs -> 1.50μs (0.599% slower)

def test_edge_palette_length_1():
    # Palette of length 1, n=1 should return the only color
    palette = ("#fff",)
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 37.8μs -> 734ns (5052% faster)
    # n=2 should raise
    with pytest.raises(ValueError):
        linear_palette(palette, 2) # 1.13μs -> 1.09μs (3.94% faster)

def test_edge_palette_length_2_various_n():
    palette = ("#1", "#2")
    # n=1 returns first
    codeflash_output = linear_palette(palette, 1) # 24.7μs -> 736ns (3256% faster)
    # n=2 returns both
    codeflash_output = linear_palette(palette, 2) # 12.3μs -> 2.97μs (312% faster)
    # n=3 raises
    with pytest.raises(ValueError):
        linear_palette(palette, 3) # 998ns -> 964ns (3.53% faster)

def test_edge_palette_with_duplicates():
    # Palette with duplicate colors
    palette = ("#abc", "#abc", "#def", "#def")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 21.2μs -> 2.57μs (726% faster)

def test_edge_palette_with_non_string():
    # Palette with non-string elements should still work (no type enforcement)
    palette = (1, 2, 3, 4)
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 19.3μs -> 2.26μs (754% faster)

def test_edge_palette_with_special_characters():
    palette = ("@", "#", "$", "%", "^")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 19.7μs -> 2.30μs (756% faster)

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

def test_large_palette_n_equals_length():
    # Large palette, n == len(palette)
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 1000); result = codeflash_output # 103μs -> 61.2μs (69.1% faster)

def test_large_palette_n_less_than_length():
    # Large palette, n < len(palette)
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 100); result = codeflash_output # 31.0μs -> 8.46μs (267% faster)
    # Should pick indices spaced by about 10
    expected = tuple(palette[int(round(i * (999/99)))] for i in range(100))

def test_large_palette_n_is_1():
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 24.0μs -> 739ns (3148% faster)

def test_large_palette_n_is_0():
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 18.5μs -> 1.83μs (910% faster)

def test_large_palette_n_greater_than_length():
    palette = tuple(f"#{i:03x}" for i in range(1000))
    with pytest.raises(ValueError):
        linear_palette(palette, 1001) # 1.36μs -> 1.29μs (5.41% faster)

def test_large_palette_with_non_hex_strings():
    palette = tuple(str(i) for i in range(1000))
    codeflash_output = linear_palette(palette, 10); result = codeflash_output # 28.8μs -> 3.06μs (841% faster)
    expected = tuple(palette[int(round(i * (999/9)))] for i in range(10))

def test_large_palette_with_duplicates():
    palette = tuple(["#abc"]*500 + ["#def"]*500)
    codeflash_output = linear_palette(palette, 10); result = codeflash_output # 23.2μs -> 2.93μs (691% faster)

# ----------------------
# Determinism Test Cases
# ----------------------

def test_determinism_same_input_same_output():
    palette = ("#a", "#b", "#c", "#d", "#e")
    codeflash_output = linear_palette(palette, 3); result1 = codeflash_output # 20.5μs -> 2.34μs (776% faster)
    codeflash_output = linear_palette(palette, 3); result2 = codeflash_output # 9.01μs -> 1.12μs (708% faster)

def test_determinism_large_palette():
    palette = tuple(f"#{i:03x}" for i in range(500))
    codeflash_output = linear_palette(palette, 50); result1 = codeflash_output # 24.5μs -> 5.22μs (369% faster)
    codeflash_output = linear_palette(palette, 50); result2 = codeflash_output # 13.1μs -> 4.05μs (222% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import math
from typing import Tuple

import numpy as np
# imports
import pytest  # used for our unit tests
from bokeh.palettes import linear_palette

Palette = Tuple[str, ...]
from bokeh.palettes import linear_palette

# unit tests

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

def test_basic_even_split():
    # 5 color palette, request 3 colors
    palette = ("#000000", "#111111", "#222222", "#333333", "#444444")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 19.5μs -> 2.31μs (742% faster)

def test_basic_full_palette():
    # n == len(palette), should return palette as-is
    palette = ("#abc", "#def", "#123")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 19.5μs -> 2.22μs (777% faster)

def test_basic_single_color():
    # n == 1, should return only the first color
    palette = ("#ff0000", "#00ff00", "#0000ff")
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 21.5μs -> 649ns (3217% faster)

def test_basic_two_colors():
    # n == 2, should return first and last color
    palette = ("#a", "#b", "#c", "#d")
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 20.9μs -> 2.62μs (699% faster)

def test_basic_three_colors_from_six():
    palette = ("#1", "#2", "#3", "#4", "#5", "#6")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 21.0μs -> 2.63μs (696% faster)

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

def test_edge_n_zero():
    # n == 0, should return empty tuple
    palette = ("#111", "#222")
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 17.4μs -> 1.98μs (779% faster)

def test_edge_palette_empty_n_zero():
    # palette empty, n == 0, should return empty tuple
    palette = ()
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 16.9μs -> 1.74μs (872% faster)

def test_edge_palette_empty_n_nonzero():
    # palette empty, n > 0, should raise ValueError
    palette = ()
    with pytest.raises(ValueError):
        linear_palette(palette, 1) # 1.27μs -> 1.30μs (2.91% slower)

def test_edge_n_greater_than_palette():
    # n > len(palette), should raise ValueError
    palette = ("#abc", "#def")
    with pytest.raises(ValueError):
        linear_palette(palette, 3) # 1.20μs -> 1.21μs (0.580% slower)

def test_edge_n_equals_palette():
    # n == len(palette), should return palette
    palette = ("#1", "#2", "#3")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 28.5μs -> 2.97μs (860% faster)

def test_edge_palette_length_one_n_one():
    # palette of length 1, n == 1, should return palette
    palette = ("#ff",)
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 23.4μs -> 666ns (3414% faster)

def test_edge_palette_length_one_n_zero():
    # palette of length 1, n == 0, should return empty tuple
    palette = ("#ff",)
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 18.1μs -> 2.05μs (782% faster)

def test_edge_palette_length_one_n_two():
    # palette of length 1, n == 2, should raise ValueError
    palette = ("#ff",)
    with pytest.raises(ValueError):
        linear_palette(palette, 2) # 1.20μs -> 1.18μs (1.87% faster)

def test_edge_palette_non_string_entries():
    # Palette contains non-string entries (should still work, as function does not validate type)
    palette = (1, 2, 3)
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 25.9μs -> 2.92μs (786% faster)

def test_edge_palette_with_duplicates():
    # Palette with duplicate colors, should still select by index
    palette = ("#a", "#a", "#b", "#b", "#c", "#c")
    codeflash_output = linear_palette(palette, 3); result = codeflash_output # 21.0μs -> 2.45μs (759% faster)

def test_edge_palette_non_hex_strings():
    # Palette with non-hex strings
    palette = ("red", "green", "blue")
    codeflash_output = linear_palette(palette, 2); result = codeflash_output # 20.2μs -> 2.26μs (795% faster)

def test_edge_palette_large_n_equals_length():
    # n == len(palette) == 1000
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 1000); result = codeflash_output # 103μs -> 60.9μs (70.3% faster)

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

def test_large_palette_large_n():
    # Large palette, n = 500
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 500); result = codeflash_output # 64.1μs -> 31.2μs (105% faster)
    # Indices should be evenly spaced; check a few
    indices = [math.floor(i) for i in np.linspace(0, 999, num=500)]
    for idx, color in zip(indices, result):
        pass

def test_large_palette_n_one():
    # Large palette, n = 1
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 1); result = codeflash_output # 23.1μs -> 775ns (2887% faster)

def test_large_palette_n_zero():
    # Large palette, n = 0
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 0); result = codeflash_output # 18.0μs -> 1.92μs (838% faster)

def test_large_palette_n_equals_length():
    # Large palette, n == len(palette)
    palette = tuple(f"#{i:03x}" for i in range(1000))
    codeflash_output = linear_palette(palette, 1000); result = codeflash_output # 102μs -> 61.0μs (67.6% faster)

def test_large_palette_n_greater_than_length():
    # Large palette, n > len(palette), should raise ValueError
    palette = tuple(f"#{i:03x}" for i in range(1000))
    with pytest.raises(ValueError):
        linear_palette(palette, 1001) # 1.39μs -> 1.35μs (3.42% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from bokeh.palettes import linear_palette
import pytest

def test_linear_palette():
    linear_palette((), 0)

def test_linear_palette_2():
    with pytest.raises(ValueError, match="Requested\\ 1\\ colors,\\ function\\ can\\ only\\ return\\ colors\\ up\\ to\\ the\\ base\\ palette's\\ length\\ \\(0\\)"):
        linear_palette((), 1)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_5f34sbte/tmpc7w0gx_9/test_concolic_coverage.py::test_linear_palette	26.6μs	2.18μs	1121%✅
codeflash_concolic_5f34sbte/tmpc7w0gx_9/test_concolic_coverage.py::test_linear_palette_2	1.30μs	1.42μs	-8.11%⚠️

To edit these changes git checkout codeflash/optimize-linear_palette-mhbgg4pu and push.