⚡️ Speed up function plasma by 14%

[codeflash-ai]

📄 14% (0.14x) speedup for plasma in src/bokeh/palettes.py

⏱️ Runtime : 1.07 milliseconds → 940 microseconds (best of 102 runs)

📝 Explanation and details

The optimized code achieves a 14% speedup by eliminating repeated math.floor() calls within the generator expression.

Key optimization: Instead of calling math.floor(i) for each floating-point index from np.linspace(), the code now:

1. Generates all floating-point indices at once with np.linspace()
2. Converts them to integers in a single vectorized operation using indices.astype(int)
3. Uses the pre-computed integer indices in a simple generator expression

Why this is faster: The original code called math.floor() for every element during tuple construction (98.6% of total time), creating significant per-element overhead. NumPy's astype(int) performs the same floor operation but vectorized in C, eliminating the Python function call overhead.

Performance benefits are most pronounced for larger palettes: Test results show the optimization provides the biggest gains when n is large (e.g., 30%+ faster for n=256, 17-25% faster for n=100), while smaller palettes (n<10) see minimal or slightly slower performance due to the overhead of the additional NumPy operations. This makes the optimization particularly valuable for applications requiring large color palettes.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
unit/bokeh/test_palettes.py::test_cmap_generator_function	28.6μs	18.3μs	56.7%✅

🌀 Generated Regression Tests and Runtime

import math

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

Plasma256 = (
    '#0C0786', '#100787', '#130689', '#15068A', '#18068B', '#1B068C', '#1D068D', '#1F058E', '#21058F', '#230590', '#250591', '#270592',
    '#290593', '#2B0594', '#2D0494', '#2F0495', '#310496', '#330497', '#340498', '#360498', '#380499', '#3A049A', '#3B039A', '#3D039B',
    '#3F039C', '#40039C', '#42039D', '#44039E', '#45039E', '#47029F', '#49029F', '#4A02A0', '#4C02A1', '#4E02A1', '#4F02A2', '#5101A2',
    '#5201A3', '#5401A3', '#5601A3', '#5701A4', '#5901A4', '#5A00A5', '#5C00A5', '#5E00A5', '#5F00A6', '#6100A6', '#6200A6', '#6400A7',
    '#6500A7', '#6700A7', '#6800A7', '#6A00A7', '#6C00A8', '#6D00A8', '#6F00A8', '#7000A8', '#7200A8', '#7300A8', '#7500A8', '#7601A8',
    '#7801A8', '#7901A8', '#7B02A8', '#7C02A7', '#7E03A7', '#7F03A7', '#8104A7', '#8204A7', '#8405A6', '#8506A6', '#8607A6', '#8807A5',
    '#8908A5', '#8B09A4', '#8C0AA4', '#8E0CA4', '#8F0DA3', '#900EA3', '#920FA2', '#9310A1', '#9511A1', '#9612A0', '#9713A0', '#99149F',
    '#9A159E', '#9B179E', '#9D189D', '#9E199C', '#9F1A9B', '#A01B9B', '#A21C9A', '#A31D99', '#A41E98', '#A51F97', '#A72197', '#A82296',
    '#A92395', '#AA2494', '#AC2593', '#AD2692', '#AE2791', '#AF2890', '#B02A8F', '#B12B8F', '#B22C8E', '#B42D8D', '#B52E8C', '#B62F8B',
    '#B7308A', '#B83289', '#B93388', '#BA3487', '#BB3586', '#BC3685', '#BD3784', '#BE3883', '#BF3982', '#C03B81', '#C13C80', '#C23D80',
    '#C33E7F', '#C43F7E', '#C5407D', '#C6417C', '#C7427B', '#C8447A', '#C94579', '#CA4678', '#CB4777', '#CC4876', '#CD4975', '#CE4A75',
    '#CF4B74', '#D04D73', '#D14E72', '#D14F71', '#D25070', '#D3516F', '#D4526E', '#D5536D', '#D6556D', '#D7566C', '#D7576B', '#D8586A',
    '#D95969', '#DA5A68', '#DB5B67', '#DC5D66', '#DC5E66', '#DD5F65', '#DE6064', '#DF6163', '#DF6262', '#E06461', '#E16560', '#E26660',
    '#E3675F', '#E3685E', '#E46A5D', '#E56B5C', '#E56C5B', '#E66D5A', '#E76E5A', '#E87059', '#E87158', '#E97257', '#EA7356', '#EA7455',
    '#EB7654', '#EC7754', '#EC7853', '#ED7952', '#ED7B51', '#EE7C50', '#EF7D4F', '#EF7E4E', '#F0804D', '#F0814D', '#F1824C', '#F2844B',
    '#F2854A', '#F38649', '#F38748', '#F48947', '#F48A47', '#F58B46', '#F58D45', '#F68E44', '#F68F43', '#F69142', '#F79241', '#F79341',
    '#F89540', '#F8963F', '#F8983E', '#F9993D', '#F99A3C', '#FA9C3B', '#FA9D3A', '#FA9F3A', '#FAA039', '#FBA238', '#FBA337', '#FBA436',
    '#FCA635', '#FCA735', '#FCA934', '#FCAA33', '#FCAC32', '#FCAD31', '#FDAF31', '#FDB030', '#FDB22F', '#FDB32E', '#FDB52D', '#FDB62D',
    '#FDB82C', '#FDB92B', '#FDBB2B', '#FDBC2A', '#FDBE29', '#FDC029', '#FDC128', '#FDC328', '#FDC427', '#FDC626', '#FCC726', '#FCC926',
    '#FCCB25', '#FCCC25', '#FCCE25', '#FBD024', '#FBD124', '#FBD324', '#FAD524', '#FAD624', '#FAD824', '#F9D924', '#F9DB24', '#F8DD24',
    '#F8DF24', '#F7E024', '#F7E225', '#F6E425', '#F6E525', '#F5E726', '#F5E926', '#F4EA26', '#F3EC26', '#F3EE26', '#F2F026', '#F2F126',
    '#F1F326', '#F0F525', '#F0F623', '#EFF821'
)
from bokeh.palettes import plasma

# unit tests

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

def test_plasma_basic_1():
    # Test for n=1 (should return the first color only)
    codeflash_output = plasma(1); result = codeflash_output # 21.2μs -> 21.9μs (3.13% slower)

def test_plasma_basic_2():
    # Test for n=2 (should return first and last color)
    codeflash_output = plasma(2); result = codeflash_output # 19.9μs -> 20.1μs (1.13% slower)

def test_plasma_basic_3():
    # Test for n=3 (should return first, middle, last color)
    codeflash_output = plasma(3); result = codeflash_output # 19.2μs -> 20.0μs (4.31% slower)

def test_plasma_basic_6():
    # Test for n=6 (example from docstring)
    codeflash_output = plasma(6); result = codeflash_output # 20.5μs -> 20.2μs (1.84% faster)

def test_plasma_basic_10():
    # Test for n=10 (should return 10 colors, evenly spaced)
    codeflash_output = plasma(10); result = codeflash_output # 20.7μs -> 20.9μs (0.795% slower)

def test_plasma_basic_full():
    # Test for n=256 (should return the full palette)
    codeflash_output = plasma(256); result = codeflash_output # 43.6μs -> 33.2μs (31.2% faster)

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

def test_plasma_zero():
    # Test for n=0 (should return empty tuple)
    codeflash_output = plasma(0); result = codeflash_output # 17.3μs -> 18.2μs (4.63% slower)


def test_plasma_overflow():
    # Test for n > 256 (should raise ValueError)
    with pytest.raises(ValueError):
        plasma(257) # 2.00μs -> 1.98μs (1.11% faster)
    with pytest.raises(ValueError):
        plasma(1000) # 903ns -> 848ns (6.49% faster)

def test_plasma_type_error():
    # Test for non-integer n (should raise TypeError from np.linspace)
    with pytest.raises(TypeError):
        plasma('a') # 1.71μs -> 1.74μs (1.84% slower)
    with pytest.raises(TypeError):
        plasma(3.5) # 4.22μs -> 3.83μs (10.2% faster)

def test_plasma_none():
    # Test for n=None (should raise TypeError)
    with pytest.raises(TypeError):
        plasma(None) # 1.46μs -> 1.51μs (3.38% slower)

def test_plasma_boundary():
    # Test for n=255 (should return 255 colors, missing only the last one)
    codeflash_output = plasma(255); result = codeflash_output # 60.2μs -> 50.5μs (19.3% faster)

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

def test_plasma_large_100():
    # Test for n=100 (large palette, should return 100 evenly spaced colors)
    codeflash_output = plasma(100); result = codeflash_output # 33.6μs -> 28.7μs (17.0% faster)

def test_plasma_large_256():
    # Test for n=256 (max allowed, should return all unique colors)
    codeflash_output = plasma(256); result = codeflash_output # 45.0μs -> 34.5μs (30.3% faster)

def test_plasma_large_500_edge():
    # Test for n=500 (should raise ValueError, as n > 256)
    with pytest.raises(ValueError):
        plasma(500) # 1.43μs -> 1.48μs (3.52% slower)

def test_plasma_performance():
    # Test performance for n=1000 (should raise ValueError, not run for performance)
    with pytest.raises(ValueError):
        plasma(1000) # 1.41μs -> 1.46μs (3.22% slower)

def test_plasma_large_spacing():
    # Test for n=256, check that each color matches the original palette
    codeflash_output = plasma(256); result = codeflash_output # 51.7μs -> 41.5μs (24.7% faster)
    for i in range(256):
        pass

# ------------------------
# Miscellaneous/Determinism Tests
# ------------------------

def test_plasma_determinism():
    # Test that repeated calls with the same n return the same result
    for n in [1, 10, 100, 256]:
        codeflash_output = plasma(n) # 82.2μs -> 67.3μs (22.2% faster)

def test_plasma_tuple_type():
    # Test that return type is always tuple
    for n in [0, 1, 2, 10, 256]:
        pass

def test_plasma_hex_format():
    # Test that all colors returned are valid hex strings
    import re
    hex_pattern = re.compile(r'^#[0-9A-F]{6})
    for n in [1, 3, 10, 256]:
        for color in plasma(n):
            pass

def test_plasma_no_mutation():
    # Test that the original Plasma256 is not mutated
    before = list(Plasma256)
    codeflash_output = plasma(10); _ = codeflash_output # 19.3μs -> 19.5μs (0.903% slower)
    after = list(Plasma256)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest  # used for our unit tests
from bokeh.palettes import plasma

Plasma256 = (
    '#0C0786', '#100787', '#130689', '#15068A', '#18068B', '#1B068C', '#1D068D', '#1F058E', '#21058F', '#230590', '#250591', '#270592',
    '#290593', '#2B0594', '#2D0494', '#2F0495', '#310496', '#330497', '#340498', '#360498', '#380499', '#3A049A', '#3B039A', '#3D039B',
    '#3F039C', '#40039C', '#42039D', '#44039E', '#45039E', '#47029F', '#49029F', '#4A02A0', '#4C02A1', '#4E02A1', '#4F02A2', '#5101A2',
    '#5201A3', '#5401A3', '#5601A3', '#5701A4', '#5901A4', '#5A00A5', '#5C00A5', '#5E00A5', '#5F00A6', '#6100A6', '#6200A6', '#6400A7',
    '#6500A7', '#6700A7', '#6800A7', '#6A00A7', '#6C00A8', '#6D00A8', '#6F00A8', '#7000A8', '#7200A8', '#7300A8', '#7500A8', '#7601A8',
    '#7801A8', '#7901A8', '#7B02A8', '#7C02A7', '#7E03A7', '#7F03A7', '#8104A7', '#8204A7', '#8405A6', '#8506A6', '#8607A6', '#8807A5',
    '#8908A5', '#8B09A4', '#8C0AA4', '#8E0CA4', '#8F0DA3', '#900EA3', '#920FA2', '#9310A1', '#9511A1', '#9612A0', '#9713A0', '#99149F',
    '#9A159E', '#9B179E', '#9D189D', '#9E199C', '#9F1A9B', '#A01B9B', '#A21C9A', '#A31D99', '#A41E98', '#A51F97', '#A72197', '#A82296',
    '#A92395', '#AA2494', '#AC2593', '#AD2692', '#AE2791', '#AF2890', '#B02A8F', '#B12B8F', '#B22C8E', '#B42D8D', '#B52E8C', '#B62F8B',
    '#B7308A', '#B83289', '#B93388', '#BA3487', '#BB3586', '#BC3685', '#BD3784', '#BE3883', '#BF3982', '#C03B81', '#C13C80', '#C23D80',
    '#C33E7F', '#C43F7E', '#C5407D', '#C6417C', '#C7427B', '#C8447A', '#C94579', '#CA4678', '#CB4777', '#CC4876', '#CD4975', '#CE4A75',
    '#CF4B74', '#D04D73', '#D14E72', '#D14F71', '#D25070', '#D3516F', '#D4526E', '#D5536D', '#D6556D', '#D7566C', '#D7576B', '#D8586A',
    '#D95969', '#DA5A68', '#DB5B67', '#DC5D66', '#DC5E66', '#DD5F65', '#DE6064', '#DF6163', '#DF6262', '#E06461', '#E16560', '#E26660',
    '#E3675F', '#E3685E', '#E46A5D', '#E56B5C', '#E56C5B', '#E66D5A', '#E76E5A', '#E87059', '#E87158', '#E97257', '#EA7356', '#EA7455',
    '#EB7654', '#EC7754', '#EC7853', '#ED7952', '#ED7B51', '#EE7C50', '#EF7D4F', '#EF7E4E', '#F0804D', '#F0814D', '#F1824C', '#F2844B',
    '#F2854A', '#F38649', '#F38748', '#F48947', '#F48A47', '#F58B46', '#F58D45', '#F68E44', '#F68F43', '#F69142', '#F79241', '#F79341',
    '#F89540', '#F8963F', '#F8983E', '#F9993D', '#F99A3C', '#FA9C3B', '#FA9D3A', '#FA9F3A', '#FAA039', '#FBA238', '#FBA337', '#FBA436',
    '#FCA635', '#FCA735', '#FCA934', '#FCAA33', '#FCAC32', '#FCAD31', '#FDAF31', '#FDB030', '#FDB22F', '#FDB32E', '#FDB52D', '#FDB62D',
    '#FDB82C', '#FDB92B', '#FDBB2B', '#FDBC2A', '#FDBE29', '#FDC029', '#FDC128', '#FDC328', '#FDC427', '#FDC626', '#FCC726', '#FCC926',
    '#FCCB25', '#FCCC25', '#FCCE25', '#FBD024', '#FBD124', '#FBD324', '#FAD524', '#FAD624', '#FAD824', '#F9D924', '#F9DB24', '#F8DD24',
    '#F8DF24', '#F7E024', '#F7E225', '#F6E425', '#F6E525', '#F5E726', '#F5E926', '#F4EA26', '#F3EC26', '#F3EE26', '#F2F026', '#F2F126',
    '#F1F326', '#F0F525', '#F0F623', '#EFF821')
from bokeh.palettes import plasma

# unit tests

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

def test_plasma_basic_single_color():
    # Test for n=1, should return the first color only
    codeflash_output = plasma(1) # 23.4μs -> 23.8μs (1.58% slower)

def test_plasma_basic_six_colors():
    # Test for n=6, should return the documented palette
    codeflash_output = plasma(6) # 22.9μs -> 23.5μs (2.24% slower)

def test_plasma_basic_full_palette():
    # Test for n=256, should return the full Plasma256 palette
    codeflash_output = plasma(256) # 44.9μs -> 34.0μs (32.1% faster)

def test_plasma_basic_typical_small_palette():
    # Test for n=3, should return three evenly spaced colors
    expected = ('#0C0786', '#B02A8F', '#EFF821')
    codeflash_output = plasma(3) # 21.0μs -> 21.3μs (1.71% slower)

def test_plasma_basic_typical_medium_palette():
    # Test for n=10, checks that output is length 10 and colors are correct
    codeflash_output = plasma(10); result = codeflash_output # 21.4μs -> 21.6μs (0.760% slower)

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

def test_plasma_edge_zero_colors():
    # Test for n=0, should return an empty tuple
    codeflash_output = plasma(0) # 17.2μs -> 17.9μs (4.18% slower)

def test_plasma_edge_n_greater_than_256():
    # Test for n=257, should raise ValueError
    with pytest.raises(ValueError):
        plasma(257) # 1.52μs -> 1.53μs (0.783% slower)


def test_plasma_edge_n_is_255():
    # Test for n=255, should return 255 colors, first and last are correct
    codeflash_output = plasma(255); result = codeflash_output # 62.3μs -> 51.7μs (20.6% faster)

def test_plasma_edge_n_is_2():
    # Test for n=2, should return first and last color
    codeflash_output = plasma(2) # 24.1μs -> 24.6μs (1.91% slower)

def test_plasma_edge_n_is_256_minus_1():
    # Test for n=255, should not include the last color twice
    codeflash_output = plasma(255); result = codeflash_output # 45.1μs -> 35.3μs (27.9% faster)

def test_plasma_edge_non_integer_input():
    # Test for non-integer input, should raise TypeError
    with pytest.raises(TypeError):
        plasma("10") # 1.83μs -> 1.86μs (2.09% slower)
    with pytest.raises(TypeError):
        plasma(10.5) # 3.10μs -> 2.75μs (13.0% faster)

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

def test_plasma_large_scale_100_colors():
    # Test for n=100, checks that output is length 100 and colors are correct
    codeflash_output = plasma(100); result = codeflash_output # 36.9μs -> 32.9μs (12.2% faster)

def test_plasma_large_scale_500_colors():
    # Test for n=500, should raise ValueError (since max is 256)
    with pytest.raises(ValueError):
        plasma(500) # 1.42μs -> 1.47μs (3.20% slower)

def test_plasma_large_scale_max_allowed():
    # Test for n=256, should return all colors, all unique
    codeflash_output = plasma(256); result = codeflash_output # 50.2μs -> 40.1μs (25.5% faster)

def test_plasma_large_scale_even_spacing():
    # Test for n=1000, should raise ValueError
    with pytest.raises(ValueError):
        plasma(1000) # 1.44μs -> 1.46μs (1.10% slower)

def test_plasma_large_scale_performance():
    # Test for n=100, ensure function returns quickly and correctly
    import time
    start = time.time()
    codeflash_output = plasma(100); result = codeflash_output # 37.2μs -> 31.4μs (18.3% faster)
    duration = time.time() - start

# ------------------ Additional Robustness Cases ------------------

def test_plasma_edge_n_is_none():
    # Test for n=None, should raise TypeError
    with pytest.raises(TypeError):
        plasma(None) # 1.67μs -> 1.75μs (4.46% slower)

def test_plasma_edge_n_is_list():
    # Test for n as a list, should raise TypeError
    with pytest.raises(TypeError):
        plasma([10]) # 1.59μs -> 1.60μs (0.875% slower)


def test_plasma_edge_palette_integrity():
    # Test that output colors are always valid hex strings
    for n in [1, 5, 10, 50, 100, 256]:
        codeflash_output = plasma(n); result = codeflash_output # 115μs -> 98.7μs (16.8% faster)
        for color in result:
            # All characters are valid hex
            int(color[1:], 16)

def test_plasma_edge_no_mutation():
    # Test that calling plasma does not mutate the global palette
    before = Plasma256[:]
    plasma(10) # 20.9μs -> 20.9μs (0.081% slower)
    after = Plasma256[:]
# 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 plasma

def test_plasma():
    plasma(0)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_5f34sbte/tmpc6qpbz9p/test_concolic_coverage.py::test_plasma	21.6μs	21.9μs	-1.47%⚠️

To edit these changes git checkout codeflash/optimize-plasma-mhbhgynk and push.