⚡️ Speed up function magma by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for magma in src/bokeh/palettes.py

⏱️ Runtime : 620 microseconds → 549 microseconds (best of 199 runs)

📝 Explanation and details

The optimization targets the expensive operation in the linear_palette function that was consuming 98.8% of the execution time. The key changes are:

What was optimized:

• Replaced the single-line generator expression that called math.floor() on each float index with a two-step process
• Split the computation into: 1) generating float indices with np.linspace(), and 2) batch-converting them to integers using NumPy's astype(int)
• Eliminated the per-element math.floor() calls by leveraging NumPy's vectorized operations

Why it's faster:

• math.floor() is expensive when called repeatedly in Python (each call has function call overhead)
• NumPy's astype(int) performs the same floor operation but vectorized across the entire array, which is much more efficient
• The optimization reduces the bottleneck from 98.8% to 53.3% of total execution time

Performance characteristics:
The optimization shows significant gains for larger palette sizes - test cases with n=100, n=255, and n=256 see 15-35% speedups, while smaller cases (n=2, n=3) see modest 3-8% improvements. This is because the vectorization benefits scale with the number of elements being processed. The 13% overall speedup reflects the mixed workload of different palette sizes in typical usage.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 5 Passed
🌀 Generated Regression Tests	✅ 27 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	30.6μs	20.5μs	49.3%✅

🌀 Generated Regression Tests and Runtime

import math
from typing import Tuple

import numpy as np
# imports
import pytest
from bokeh.palettes import magma

Palette = Tuple[str, ...]
Magma256 = (
    '#000003', '#000004', '#000006', '#010007', '#010109', '#01010B', '#02020D', '#02020F', '#030311', '#040313', '#040415', '#050417',
    '#060519', '#07051B', '#08061D', '#09071F', '#0A0722', '#0B0824', '#0C0926', '#0D0A28', '#0E0A2A', '#0F0B2C', '#100C2F', '#110C31',
    '#120D33', '#140D35', '#150E38', '#160E3A', '#170F3C', '#180F3F', '#1A1041', '#1B1044', '#1C1046', '#1E1049', '#1F114B', '#20114D',
    '#221150', '#231152', '#251155', '#261157', '#281159', '#2A115C', '#2B115E', '#2D1060', '#2F1062', '#301065', '#321067', '#341068',
    '#350F6A', '#370F6C', '#390F6E', '#3B0F6F', '#3C0F71', '#3E0F72', '#400F73', '#420F74', '#430F75', '#450F76', '#470F77', '#481078',
    '#4A1079', '#4B1079', '#4D117A', '#4F117B', '#50127B', '#52127C', '#53137C', '#55137D', '#57147D', '#58157E', '#5A157E', '#5B167E',
    '#5D177E', '#5E177F', '#60187F', '#61187F', '#63197F', '#651A80', '#661A80', '#681B80', '#691C80', '#6B1C80', '#6C1D80', '#6E1E81',
    '#6F1E81', '#711F81', '#731F81', '#742081', '#762181', '#772181', '#792281', '#7A2281', '#7C2381', '#7E2481', '#7F2481', '#812581',
    '#822581', '#842681', '#852681', '#872781', '#892881', '#8A2881', '#8C2980', '#8D2980', '#8F2A80', '#912A80', '#922B80', '#942B80',
    '#952C80', '#972C7F', '#992D7F', '#9A2D7F', '#9C2E7F', '#9E2E7E', '#9F2F7E', '#A12F7E', '#A3307E', '#A4307D', '#A6317D', '#A7317D',
    '#A9327C', '#AB337C', '#AC337B', '#AE347B', '#B0347B', '#B1357A', '#B3357A', '#B53679', '#B63679', '#B83778', '#B93778', '#BB3877',
    '#BD3977', '#BE3976', '#C03A75', '#C23A75', '#C33B74', '#C53C74', '#C63C73', '#C83D72', '#CA3E72', '#CB3E71', '#CD3F70', '#CE4070',
    '#D0416F', '#D1426E', '#D3426D', '#D4436D', '#D6446C', '#D7456B', '#D9466A', '#DA4769', '#DC4869', '#DD4968', '#DE4A67', '#E04B66',
    '#E14C66', '#E24D65', '#E44E64', '#E55063', '#E65162', '#E75262', '#E85461', '#EA5560', '#EB5660', '#EC585F', '#ED595F', '#EE5B5E',
    '#EE5D5D', '#EF5E5D', '#F0605D', '#F1615C', '#F2635C', '#F3655C', '#F3675B', '#F4685B', '#F56A5B', '#F56C5B', '#F66E5B', '#F6705B',
    '#F7715B', '#F7735C', '#F8755C', '#F8775C', '#F9795C', '#F97B5D', '#F97D5D', '#FA7F5E', '#FA805E', '#FA825F', '#FB8460', '#FB8660',
    '#FB8861', '#FB8A62', '#FC8C63', '#FC8E63', '#FC9064', '#FC9265', '#FC9366', '#FD9567', '#FD9768', '#FD9969', '#FD9B6A', '#FD9D6B',
    '#FD9F6C', '#FDA16E', '#FDA26F', '#FDA470', '#FEA671', '#FEA873', '#FEAA74', '#FEAC75', '#FEAE76', '#FEAF78', '#FEB179', '#FEB37B',
    '#FEB57C', '#FEB77D', '#FEB97F', '#FEBB80', '#FEBC82', '#FEBE83', '#FEC085', '#FEC286', '#FEC488', '#FEC689', '#FEC78B', '#FEC98D',
    '#FECB8E', '#FDCD90', '#FDCF92', '#FDD193', '#FDD295', '#FDD497', '#FDD698', '#FDD89A', '#FDDA9C', '#FDDC9D', '#FDDD9F', '#FDDFA1',
    '#FDE1A3', '#FCE3A5', '#FCE5A6', '#FCE6A8', '#FCE8AA', '#FCEAAC', '#FCECAE', '#FCEEB0', '#FCF0B1', '#FCF1B3', '#FCF3B5', '#FCF5B7',
    '#FBF7B9', '#FBF9BB', '#FBFABD', '#FBFCBF'
)
from bokeh.palettes import magma

# --------------------------
# Unit tests for magma
# --------------------------

# 1. Basic Test Cases

def test_magma_single_color():
    # Should return the first color when n=1
    codeflash_output = magma(1) # 27.2μs -> 27.0μs (0.778% faster)

def test_magma_two_colors():
    # Should return first and last color for n=2
    codeflash_output = magma(2) # 21.7μs -> 21.0μs (3.30% faster)

def test_magma_three_colors():
    # Should return first, middle, last
    codeflash_output = magma(3) # 22.0μs -> 20.5μs (7.64% faster)

def test_magma_six_colors():
    # Test example from docstring
    codeflash_output = magma(6) # 21.7μs -> 21.6μs (0.240% faster)

def test_magma_typical_small_n():
    # Should return correct palette for n=5
    codeflash_output = magma(5); result = codeflash_output # 21.2μs -> 20.2μs (4.77% faster)
    expected = (
        Magma256[0],
        Magma256[64],
        Magma256[128],
        Magma256[192],
        Magma256[255]
    )

def test_magma_typical_medium_n():
    # Should return correct palette for n=10
    codeflash_output = magma(10); result = codeflash_output # 21.6μs -> 21.5μs (0.405% faster)
    expected = tuple(Magma256[math.floor(i)] for i in np.linspace(0, 255, 10))

# 2. Edge Test Cases

def test_magma_zero_colors():
    # Should return an empty tuple if n=0
    codeflash_output = magma(0) # 17.1μs -> 17.0μs (0.996% faster)

def test_magma_max_palette():
    # Should return the entire palette for n=256
    codeflash_output = magma(256) # 43.5μs -> 33.7μs (29.1% faster)

def test_magma_n_greater_than_palette():
    # Should raise ValueError if n > 256
    with pytest.raises(ValueError):
        magma(257) # 1.55μs -> 1.73μs (10.4% slower)


def test_magma_non_integer_n():
    # Should raise TypeError if n is not integer
    with pytest.raises(TypeError):
        magma(3.5) # 4.49μs -> 4.31μs (4.06% faster)

def test_magma_string_input():
    # Should raise TypeError if n is a string
    with pytest.raises(TypeError):
        magma("10") # 1.83μs -> 1.87μs (2.56% slower)

def test_magma_none_input():
    # Should raise TypeError if n is None
    with pytest.raises(TypeError):
        magma(None) # 1.56μs -> 1.65μs (5.75% slower)

def test_magma_boundary_indices():
    # Check that indices are correct for n=4
    codeflash_output = magma(4); result = codeflash_output # 36.7μs -> 37.3μs (1.66% slower)
    # Should be indices 0, 85, 170, 255
    expected = (Magma256[0], Magma256[85], Magma256[170], Magma256[255])

def test_magma_all_colors_are_hex():
    # All returned colors should be hex strings starting with '#'
    codeflash_output = magma(20); palette = codeflash_output # 26.8μs -> 25.9μs (3.27% faster)

# 3. Large Scale Test Cases

def test_magma_large_n_100():
    # Should return palette of length 100
    codeflash_output = magma(100); result = codeflash_output # 32.7μs -> 28.5μs (14.7% faster)

def test_magma_large_n_255():
    # Should return palette of length 255
    codeflash_output = magma(255); result = codeflash_output # 44.5μs -> 34.4μs (29.2% faster)

def test_magma_large_n_256():
    # Should return all colors, and each color is unique
    codeflash_output = magma(256); result = codeflash_output # 43.8μs -> 33.1μs (32.5% faster)

def test_magma_performance_large_n():
    # Should run efficiently for n=999 (max allowed is 256, so expect ValueError)
    with pytest.raises(ValueError):
        magma(999) # 1.51μs -> 1.68μs (10.0% slower)

def test_magma_progression():
    # For n=10, colors should progress monotonically through the palette
    codeflash_output = magma(10); result = codeflash_output # 27.7μs -> 27.4μs (1.01% faster)
    indices = [math.floor(i) for i in np.linspace(0, 255, 10)]
    for idx, color in zip(indices, result):
        pass

# Additional edge: test that consecutive calls are independent
def test_magma_statelessness():
    # Calling magma with different n should not affect each other
    codeflash_output = magma(5); p1 = codeflash_output # 20.4μs -> 20.4μs (0.157% faster)
    codeflash_output = magma(10); p2 = codeflash_output # 9.36μs -> 9.52μs (1.67% slower)

# Additional: test that repeated colors only occur when n > len(set(Magma256))
def test_magma_no_duplicates_for_n_leq_256():
    # For n <= 256, all colors should be unique
    for n in [1, 2, 10, 100, 256]:
        codeflash_output = magma(n); result = codeflash_output # 82.8μs -> 69.8μs (18.6% faster)

# Additional: test that for n=256, all indices are present
def test_magma_indices_for_n_256():
    codeflash_output = magma(256); result = codeflash_output # 41.3μs -> 30.6μs (34.7% faster)
    for i, color in enumerate(result):
        pass
# 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 magma

def test_magma():
    magma(0)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_5f34sbte/tmpiq527na1/test_concolic_coverage.py::test_magma	16.7μs	17.4μs	-3.73%⚠️

To edit these changes git checkout codeflash/optimize-magma-mhbh4wu3 and push.