⚡️ Speed up function _set_cmap by 4,516%

[codeflash-ai]

📄 4,516% (45.16x) speedup for _set_cmap in optuna/visualization/matplotlib/_contour.py

⏱️ Runtime : 26.7 milliseconds → 579 microseconds (best of 250 runs)

📝 Explanation and details

The optimization pre-computes the colormap objects at module level instead of calling plt.get_cmap() on every function invocation.

Key changes:

• Added module-level variables _blues_r_cmap and _blues_cmap that cache the colormap objects
• Replaced the plt.get_cmap(cmap) call with direct object references

Why this speeds up the code:
The line profiler shows that plt.get_cmap(cmap) was consuming 97.2% of the function's runtime (40.88ms out of 42.04ms total). This call involves matplotlib's internal colormap lookup and object creation/retrieval mechanisms, which are expensive operations. By pre-computing these objects once at import time and storing them as module constants, we eliminate this overhead entirely on each function call.

Performance impact:

• 46x speedup (4516% improvement) from 26.7ms to 579μs
• All test cases show consistent 30-50x improvements, demonstrating the optimization works across different input types
• The bulk test cases (test_set_cmap_many_calls_performance, test_set_cmap_with_varied_inputs) show particularly strong gains (52x speedup), indicating this optimization scales well with high call frequency

Workload benefits:
This optimization is especially valuable for visualization code that generates multiple plots or updates plots frequently, as colormap selection is typically called repeatedly during rendering operations. The consistent speedup across all test cases suggests any workload calling this function will benefit significantly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from optuna.visualization.matplotlib._contour import _set_cmap


# Simulate plt.get_cmap and Colormap for isolated testing (since we can't import matplotlib here)
class DummyColormap:
    def __init__(self, name):
        self.name = name

def dummy_get_cmap(name):
    # Simulate matplotlib's get_cmap behavior
    if name not in ["Blues", "Blues_r"]:
        raise ValueError(f"Unknown colormap: {name}")
    return DummyColormap(name)

# Patch plt for the test context
class DummyPlt:
    get_cmap = staticmethod(dummy_get_cmap)
from optuna.visualization.matplotlib._contour import _set_cmap

# unit tests

# 1. Basic Test Cases

def test_set_cmap_false_returns_blues_r():
    # Test that False returns the reversed colormap
    codeflash_output = _set_cmap(False); result = codeflash_output # 14.5μs -> 395ns (3579% faster)

def test_set_cmap_true_returns_blues():
    # Test that True returns the normal colormap
    codeflash_output = _set_cmap(True); result = codeflash_output # 12.0μs -> 309ns (3772% faster)

def test_set_cmap_bool_true_and_false():
    # Test both True and False, ensure mapping is correct
    codeflash_output = _set_cmap(True).name # 12.3μs -> 329ns (3636% faster)
    codeflash_output = _set_cmap(False).name # 7.56μs -> 187ns (3943% faster)

# 2. Edge Test Cases

def test_set_cmap_with_non_bool_true_like():
    # Test with a truthy non-bool value (should treat as True)
    codeflash_output = _set_cmap(1); result = codeflash_output # 11.3μs -> 338ns (3230% faster)

def test_set_cmap_with_non_bool_false_like():
    # Test with a falsy non-bool value (should treat as False)
    codeflash_output = _set_cmap(0); result = codeflash_output # 11.3μs -> 300ns (3663% faster)

def test_set_cmap_with_none():
    # None is falsy, should return "Blues_r"
    codeflash_output = _set_cmap(None); result = codeflash_output # 11.6μs -> 366ns (3083% faster)

def test_set_cmap_with_custom_bool_object():
    # Custom object with __bool__ returning True
    class Truthy:
        def __bool__(self): return True
    codeflash_output = _set_cmap(Truthy()).name # 12.7μs -> 779ns (1531% faster)

    # Custom object with __bool__ returning False
    class Falsy:
        def __bool__(self): return False
    codeflash_output = _set_cmap(Falsy()).name # 8.00μs -> 278ns (2777% faster)

def test_set_cmap_with_string_true_false():
    # Non-empty string is truthy, empty string is falsy
    codeflash_output = _set_cmap("True").name # 11.4μs -> 349ns (3168% faster)
    codeflash_output = _set_cmap("").name # 7.61μs -> 202ns (3666% faster)

def test_set_cmap_with_large_integer():
    # Large positive integer is truthy
    codeflash_output = _set_cmap(999999).name # 11.1μs -> 341ns (3153% faster)
    # Large negative integer is also truthy
    codeflash_output = _set_cmap(-123456).name # 7.15μs -> 193ns (3606% faster)

def test_set_cmap_with_empty_list():
    # Empty list is falsy
    codeflash_output = _set_cmap([]).name # 10.7μs -> 337ns (3075% faster)

def test_set_cmap_with_non_empty_list():
    # Non-empty list is truthy
    codeflash_output = _set_cmap([1, 2, 3]).name # 11.1μs -> 335ns (3221% faster)


def test_set_cmap_many_calls_performance():
    # Make sure function can handle many calls without error or state leak
    for i in range(1000):
        val = (i % 2 == 0)
        codeflash_output = _set_cmap(val); cmap = codeflash_output # 6.12ms -> 114μs (5234% faster)
        expected = "Blues_r" if not val else "Blues"

def test_set_cmap_with_varied_inputs():
    # Feed a large variety of truthy/falsy values
    inputs = [True, False, 1, 0, -1, [], [1], {}, {1:2}, "", "x", None, set(), {1}, (), (1,), 3.14, 0.0]
    expected_names = []
    for inp in inputs:
        expected = "Blues" if inp else "Blues_r"
        codeflash_output = _set_cmap(inp); cmap = codeflash_output # 120μs -> 2.82μs (4194% faster)
        expected_names.append(expected)

def test_set_cmap_with_bulk_custom_objects():
    # Create 1000 custom objects, half truthy, half falsy
    class T:
        def __init__(self, b): self.b = b
        def __bool__(self): return self.b
    objs = [T(i % 2 == 0) for i in range(1000)]
    for i, obj in enumerate(objs):
        expected = "Blues" if obj else "Blues_r"
        codeflash_output = _set_cmap(obj).name # 6.17ms -> 171μs (3497% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from optuna.visualization.matplotlib._contour import _set_cmap


# Minimal mock of plt.get_cmap for testing, since we can't import matplotlib in this environment.
class DummyColormap:
    def __init__(self, name):
        self.name = name

class DummyPlt:
    def __init__(self):
        self._colormaps = {
            "Blues": DummyColormap("Blues"),
            "Blues_r": DummyColormap("Blues_r"),
        }
    def get_cmap(self, name):
        if name not in self._colormaps:
            raise ValueError(f"Colormap {name} not found")
        return self._colormaps[name]

# Patch plt in the function's namespace for testing
plt = DummyPlt()
from optuna.visualization.matplotlib._contour import _set_cmap

# unit tests

# 1. Basic Test Cases

def test_set_cmap_returns_blues_r_when_reverse_scale_false():
    """Test that _set_cmap returns 'Blues_r' colormap when reverse_scale is False."""
    codeflash_output = _set_cmap(False); cmap = codeflash_output # 14.2μs -> 383ns (3600% faster)

def test_set_cmap_returns_blues_when_reverse_scale_true():
    """Test that _set_cmap returns 'Blues' colormap when reverse_scale is True."""
    codeflash_output = _set_cmap(True); cmap = codeflash_output # 11.6μs -> 327ns (3435% faster)

# 2. Edge Test Cases

def test_set_cmap_with_non_bool_true_value():
    """Test that _set_cmap treats truthy non-bool values as True (should return 'Blues')."""
    codeflash_output = _set_cmap(1); cmap = codeflash_output # 11.4μs -> 360ns (3075% faster)

def test_set_cmap_with_non_bool_false_value():
    """Test that _set_cmap treats falsy non-bool values as False (should return 'Blues_r')."""
    codeflash_output = _set_cmap(0); cmap = codeflash_output # 11.6μs -> 328ns (3427% faster)

def test_set_cmap_with_none():
    """Test that _set_cmap treats None as False (should return 'Blues_r')."""
    codeflash_output = _set_cmap(None); cmap = codeflash_output # 11.9μs -> 312ns (3723% faster)

def test_set_cmap_with_unexpected_type_string():
    """Test that _set_cmap treats non-empty string as True (should return 'Blues')."""
    codeflash_output = _set_cmap("nonempty"); cmap = codeflash_output # 11.7μs -> 328ns (3452% faster)

def test_set_cmap_with_empty_string():
    """Test that _set_cmap treats empty string as False (should return 'Blues_r')."""
    codeflash_output = _set_cmap(""); cmap = codeflash_output # 11.8μs -> 302ns (3817% faster)

def test_set_cmap_with_list():
    """Test that _set_cmap treats non-empty list as True (should return 'Blues')."""
    codeflash_output = _set_cmap([1]); cmap = codeflash_output # 12.2μs -> 376ns (3144% faster)

def test_set_cmap_with_empty_list():
    """Test that _set_cmap treats empty list as False (should return 'Blues_r')."""
    codeflash_output = _set_cmap([]); cmap = codeflash_output # 11.8μs -> 335ns (3410% faster)


def test_set_cmap_many_true_and_false():
    """Test _set_cmap repeatedly with many True/False values to check for consistency and performance."""
    for i in range(500):
        codeflash_output = _set_cmap(True); cmap1 = codeflash_output # 3.05ms -> 56.9μs (5265% faster)
        codeflash_output = _set_cmap(False); cmap2 = codeflash_output # 3.05ms -> 55.6μs (5383% faster)

def test_set_cmap_with_various_truthy_and_falsy_large():
    """Test _set_cmap with many different truthy/falsy values in a large batch."""
    truthy_values = [1, "yes", [1], {"a": 1}, (1,), True]
    falsy_values = [0, "", [], {}, (), None, False]
    for i in range(100):
        for val in truthy_values:
            codeflash_output = _set_cmap(val); cmap = codeflash_output
        for val in falsy_values:
            codeflash_output = _set_cmap(val); cmap = codeflash_output
# 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-_set_cmap-mhoaygqj and push.