⚡️ Speed up function _is_log_scale by 225%

[codeflash-ai]

📄 225% (2.25x) speedup for _is_log_scale in optuna/visualization/matplotlib/_utils.py

⏱️ Runtime : 2.07 milliseconds → 638 microseconds (best of 145 runs)

📝 Explanation and details

The optimization replaces isinstance(dist, (FloatDistribution, IntDistribution)) with a more efficient type checking approach. Instead of using isinstance() with a tuple of types, the code now:

1. Caches the type once: dist_type = type(dist) stores the object's type in a local variable
2. Uses identity comparison: dist_type is FloatDistribution or dist_type is IntDistribution performs exact type matching

Why this is faster:

• isinstance() with a tuple requires iterating through each type in the tuple and performing inheritance checks
• type() + identity comparison (is) is a single pointer comparison operation, which is much faster
• The is operator checks object identity directly, avoiding the overhead of inheritance hierarchy traversal

Performance characteristics:
The optimization shows consistent 60-270% speedups across test cases, with particularly strong gains on:

• Large-scale tests with many trials (263-269% faster)
• Cases with non-standard distribution types (980-1700% faster)
• Single trial scenarios (80-133% faster)

The speedup is most pronounced when the function processes many trials or encounters non-FloatDistribution/IntDistribution objects, as the isinstance() overhead compounds with scale while the optimized version maintains constant-time type checking.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from types import SimpleNamespace

# imports
import pytest
from optuna.visualization.matplotlib._utils import _is_log_scale


# --- Function to test ---
# Mocked minimal versions of the Optuna classes for testing purposes.
class FloatDistribution:
    def __init__(self, log=False):
        self.log = log

class IntDistribution:
    def __init__(self, log=False):
        self.log = log

class FrozenTrial:
    def __init__(self, params=None, distributions=None):
        self.params = params or {}
        self.distributions = distributions or {}
from optuna.visualization.matplotlib._utils import _is_log_scale

# --- Unit tests ---

# 1. Basic Test Cases

def test_empty_trials_returns_false():
    # No trials at all
    codeflash_output = _is_log_scale([], "x") # 428ns -> 470ns (8.94% slower)

def test_no_trials_with_param_returns_false():
    # Trials exist but none have the parameter
    t1 = FrozenTrial(params={}, distributions={})
    t2 = FrozenTrial(params={}, distributions={})
    codeflash_output = _is_log_scale([t1, t2], "x") # 723ns -> 737ns (1.90% slower)

def test_single_trial_param_not_log_returns_false():
    # Single trial with param, but log=False
    t = FrozenTrial(params={"x": 1.0}, distributions={"x": FloatDistribution(log=False)})
    codeflash_output = _is_log_scale([t], "x") # 1.99μs -> 904ns (120% faster)

def test_single_trial_param_log_returns_true():
    # Single trial with param, log=True
    t = FrozenTrial(params={"x": 1.0}, distributions={"x": FloatDistribution(log=True)})
    codeflash_output = _is_log_scale([t], "x") # 1.57μs -> 858ns (83.6% faster)

def test_multiple_trials_first_with_log_returns_true():
    # Multiple trials, first with log=True
    t1 = FrozenTrial(params={"x": 1}, distributions={"x": IntDistribution(log=True)})
    t2 = FrozenTrial(params={"x": 2}, distributions={"x": IntDistribution(log=False)})
    codeflash_output = _is_log_scale([t1, t2], "x") # 2.49μs -> 1.07μs (133% faster)

def test_multiple_trials_later_with_log_returns_true():
    # Multiple trials, log=True in a later trial
    t1 = FrozenTrial(params={"x": 1}, distributions={"x": IntDistribution(log=False)})
    t2 = FrozenTrial(params={"x": 2}, distributions={"x": IntDistribution(log=True)})
    codeflash_output = _is_log_scale([t1, t2], "x") # 1.89μs -> 991ns (91.0% faster)

def test_multiple_trials_all_without_log_returns_false():
    # Multiple trials, none with log=True
    t1 = FrozenTrial(params={"x": 1}, distributions={"x": IntDistribution(log=False)})
    t2 = FrozenTrial(params={"x": 2}, distributions={"x": IntDistribution(log=False)})
    codeflash_output = _is_log_scale([t1, t2], "x") # 1.95μs -> 1.02μs (90.4% faster)

# 2. Edge Test Cases

def test_param_in_some_trials_only():
    # Only some trials have the param, one with log=True
    t1 = FrozenTrial(params={}, distributions={})
    t2 = FrozenTrial(params={"x": 2}, distributions={"x": FloatDistribution(log=True)})
    t3 = FrozenTrial(params={}, distributions={})
    codeflash_output = _is_log_scale([t1, t2, t3], "x") # 1.67μs -> 988ns (68.5% faster)

def test_param_in_some_trials_none_with_log():
    # Only some trials have the param, none with log=True
    t1 = FrozenTrial(params={}, distributions={})
    t2 = FrozenTrial(params={"x": 2}, distributions={"x": FloatDistribution(log=False)})
    t3 = FrozenTrial(params={}, distributions={})
    codeflash_output = _is_log_scale([t1, t2, t3], "x") # 1.57μs -> 944ns (66.2% faster)

def test_param_with_non_floatint_distribution():
    # The distribution is not FloatDistribution or IntDistribution
    class DummyDistribution:
        def __init__(self, log=True):
            self.log = log
    t = FrozenTrial(params={"x": 1}, distributions={"x": DummyDistribution(log=True)})
    # Should ignore DummyDistribution and return False
    codeflash_output = _is_log_scale([t], "x") # 15.7μs -> 870ns (1703% faster)

def test_param_with_missing_distribution_key():
    # The param is present in params but missing in distributions
    t = FrozenTrial(params={"x": 1}, distributions={})
    # Should raise KeyError
    with pytest.raises(KeyError):
        _is_log_scale([t], "x") # 1.20μs -> 1.12μs (7.34% faster)



def test_param_with_non_bool_log():
    # The 'log' attribute is not a boolean
    t = FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log="yes")})
    # Should treat any truthy value as True
    codeflash_output = _is_log_scale([t], "x") # 2.10μs -> 946ns (122% faster)

def test_param_with_falsey_non_bool_log():
    # The 'log' attribute is a falsey non-bool value
    t = FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log="")})
    codeflash_output = _is_log_scale([t], "x") # 1.58μs -> 878ns (80.1% faster)

def test_param_with_multiple_params_in_trial():
    # Trial has multiple params, only one is log=True
    t = FrozenTrial(
        params={"x": 1, "y": 2},
        distributions={"x": FloatDistribution(log=True), "y": IntDistribution(log=False)}
    )
    codeflash_output = _is_log_scale([t], "x") # 1.57μs -> 845ns (85.4% faster)
    codeflash_output = _is_log_scale([t], "y") # 1.01μs -> 322ns (212% faster)

def test_param_with_case_sensitive_param_name():
    # Param name is case sensitive
    t = FrozenTrial(params={"X": 1}, distributions={"X": FloatDistribution(log=True)})
    codeflash_output = _is_log_scale([t], "x") # 583ns -> 571ns (2.10% faster)
    codeflash_output = _is_log_scale([t], "X") # 1.22μs -> 577ns (112% faster)

# 3. Large Scale Test Cases

def test_large_number_of_trials_all_false():
    # 1000 trials, none with log=True
    trials = [FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)}) for i in range(1000)]
    codeflash_output = _is_log_scale(trials, "x") # 270μs -> 74.4μs (264% faster)

def test_large_number_of_trials_one_true():
    # 999 trials with log=False, 1 with log=True
    trials = [FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)}) for i in range(999)]
    trials.append(FrozenTrial(params={"x": 1000}, distributions={"x": FloatDistribution(log=True)}))
    codeflash_output = _is_log_scale(trials, "x") # 266μs -> 71.3μs (273% faster)

def test_large_number_of_trials_param_missing_in_most():
    # 1000 trials, only last has param with log=True
    trials = [FrozenTrial(params={}, distributions={}) for _ in range(999)]
    trials.append(FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log=True)}))
    codeflash_output = _is_log_scale(trials, "x") # 25.6μs -> 24.8μs (3.52% faster)

def test_large_number_of_trials_param_missing_everywhere():
    # 1000 trials, none have the param
    trials = [FrozenTrial(params={}, distributions={}) for _ in range(1000)]
    codeflash_output = _is_log_scale(trials, "x") # 24.7μs -> 24.8μs (0.311% slower)

def test_large_number_of_trials_mixed_types():
    # 500 FloatDistribution, 500 IntDistribution, all log=False
    trials = []
    for i in range(500):
        trials.append(FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)}))
    for i in range(500):
        trials.append(FrozenTrial(params={"x": i}, distributions={"x": IntDistribution(log=False)}))
    codeflash_output = _is_log_scale(trials, "x") # 264μs -> 73.5μs (259% faster)

def test_large_number_of_trials_mixed_types_with_one_log_true():
    # 999 trials with log=False, last trial IntDistribution with log=True
    trials = []
    for i in range(500):
        trials.append(FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)}))
    for i in range(499):
        trials.append(FrozenTrial(params={"x": i}, distributions={"x": IntDistribution(log=False)}))
    trials.append(FrozenTrial(params={"x": 1000}, distributions={"x": IntDistribution(log=True)}))
    codeflash_output = _is_log_scale(trials, "x") # 260μs -> 72.1μs (262% 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
from optuna.visualization.matplotlib._utils import _is_log_scale


# function to test
# Simulated minimal versions of optuna classes for testing
class FloatDistribution:
    def __init__(self, log=False):
        self.log = log

class IntDistribution:
    def __init__(self, log=False):
        self.log = log

class FrozenTrial:
    def __init__(self, params, distributions):
        self.params = params  # dict: param_name -> value
        self.distributions = distributions  # dict: param_name -> distribution instance
from optuna.visualization.matplotlib._utils import _is_log_scale

# unit tests

# ========== Basic Test Cases ==========

def test_empty_trials_returns_false():
    # No trials: always return False
    codeflash_output = _is_log_scale([], "x") # 462ns -> 464ns (0.431% slower)

def test_param_not_in_any_trial_returns_false():
    # Trials exist, but param is not present in any
    trials = [
        FrozenTrial(params={"a": 1}, distributions={"a": FloatDistribution(log=False)}),
        FrozenTrial(params={"b": 2}, distributions={"b": IntDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 671ns -> 691ns (2.89% slower)

def test_param_in_trial_with_floatdistribution_log_false_returns_false():
    # Param present, but log is False
    trials = [
        FrozenTrial(params={"x": 1.5}, distributions={"x": FloatDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 2.03μs -> 870ns (133% faster)

def test_param_in_trial_with_intdistribution_log_false_returns_false():
    # Param present, but log is False (IntDistribution)
    trials = [
        FrozenTrial(params={"x": 2}, distributions={"x": IntDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.81μs -> 816ns (122% faster)

def test_param_in_trial_with_floatdistribution_log_true_returns_true():
    # Param present, FloatDistribution, log=True
    trials = [
        FrozenTrial(params={"x": 2.0}, distributions={"x": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.53μs -> 823ns (86.4% faster)

def test_param_in_trial_with_intdistribution_log_true_returns_true():
    # Param present, IntDistribution, log=True
    trials = [
        FrozenTrial(params={"x": 3}, distributions={"x": IntDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.41μs -> 767ns (84.1% faster)

def test_param_in_multiple_trials_first_false_second_true_returns_true():
    # Multiple trials, first has log=False, second has log=True
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log=False)}),
        FrozenTrial(params={"x": 2}, distributions={"x": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 2.03μs -> 1.11μs (83.0% faster)

def test_param_in_multiple_trials_all_false_returns_false():
    # Multiple trials, all have log=False
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log=False)}),
        FrozenTrial(params={"x": 2}, distributions={"x": IntDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.99μs -> 1.07μs (86.4% faster)

def test_param_in_some_trials_only_returns_true_if_any_log_true():
    # Some trials have the param, some don't; only one has log=True
    trials = [
        FrozenTrial(params={"y": 1}, distributions={"y": FloatDistribution(log=False)}),
        FrozenTrial(params={"x": 2}, distributions={"x": FloatDistribution(log=True)}),
        FrozenTrial(params={"x": 3}, distributions={"x": FloatDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.99μs -> 1.12μs (78.3% faster)

def test_param_in_some_trials_only_returns_false_if_none_log_true():
    # Some trials have the param, some don't; none have log=True
    trials = [
        FrozenTrial(params={"y": 1}, distributions={"y": FloatDistribution(log=False)}),
        FrozenTrial(params={"x": 2}, distributions={"x": FloatDistribution(log=False)}),
        FrozenTrial(params={"x": 3}, distributions={"x": FloatDistribution(log=False)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 2.01μs -> 1.12μs (79.0% faster)

# ========== Edge Test Cases ==========

def test_param_in_trial_with_non_float_or_int_distribution_returns_false():
    # Distribution is not FloatDistribution or IntDistribution
    class DummyDistribution:
        def __init__(self, log=True):
            self.log = log
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": DummyDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 14.1μs -> 817ns (1626% faster)

def test_param_in_trial_with_distribution_missing_log_attribute():
    # Distribution does not have a 'log' attribute
    class NoLogDistribution:
        pass
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": NoLogDistribution()}),
    ]
    # Should not raise, should return False
    codeflash_output = _is_log_scale(trials, "x") # 9.92μs -> 745ns (1231% faster)

def test_param_in_trial_with_distribution_log_is_none():
    # Distribution's log attribute is None
    class WeirdDistribution(FloatDistribution):
        def __init__(self):
            self.log = None
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": WeirdDistribution()}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 8.86μs -> 780ns (1036% faster)

def test_param_in_trial_with_distribution_log_is_non_bool():
    # Distribution's log attribute is a non-bool truthy value
    class WeirdDistribution(FloatDistribution):
        def __init__(self):
            self.log = 1  # truthy, should be treated as True
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": WeirdDistribution()}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 7.98μs -> 739ns (980% faster)

def test_param_in_trial_with_distribution_log_is_non_bool_falsey():
    # Distribution's log attribute is a non-bool falsey value
    class WeirdDistribution(FloatDistribution):
        def __init__(self):
            self.log = 0  # falsey, should be treated as False
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": WeirdDistribution()}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 8.00μs -> 787ns (917% faster)

def test_param_is_empty_string():
    # Param is an empty string
    trials = [
        FrozenTrial(params={"": 1}, distributions={"": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "") # 1.49μs -> 765ns (94.1% faster)

def test_param_is_none():
    # Param is None (should not match any param)
    trials = [
        FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, None) # 623ns -> 644ns (3.26% slower)

def test_trial_with_empty_params_dict():
    # Trial has empty params dict
    trials = [
        FrozenTrial(params={}, distributions={}),
        FrozenTrial(params={"x": 1}, distributions={"x": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 1.67μs -> 916ns (81.9% faster)

def test_trial_with_params_but_missing_distribution():
    # Param present in params, but missing in distributions
    trials = [
        FrozenTrial(params={"x": 1}, distributions={}),
    ]
    with pytest.raises(KeyError):
        _is_log_scale(trials, "x") # 1.14μs -> 1.17μs (2.74% slower)

def test_trial_with_distribution_but_missing_param():
    # Distribution present, but param not in params
    trials = [
        FrozenTrial(params={}, distributions={"x": FloatDistribution(log=True)}),
    ]
    codeflash_output = _is_log_scale(trials, "x") # 600ns -> 680ns (11.8% slower)

# ========== Large Scale Test Cases ==========

def test_large_number_of_trials_all_false():
    # 1000 trials, none with log=True
    trials = [
        FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)})
        for i in range(1000)
    ]
    codeflash_output = _is_log_scale(trials, "x") # 265μs -> 73.1μs (263% faster)

def test_large_number_of_trials_one_true_early():
    # 999 trials with log=False, 1st with log=True
    trials = [
        FrozenTrial(params={"x": 0}, distributions={"x": FloatDistribution(log=True)})
    ] + [
        FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)})
        for i in range(1, 1000)
    ]
    codeflash_output = _is_log_scale(trials, "x") # 264μs -> 72.0μs (267% faster)

def test_large_number_of_trials_one_true_late():
    # 999 trials with log=False, last with log=True
    trials = [
        FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)})
        for i in range(999)
    ] + [
        FrozenTrial(params={"x": 999}, distributions={"x": FloatDistribution(log=True)})
    ]
    codeflash_output = _is_log_scale(trials, "x") # 263μs -> 71.3μs (269% faster)

def test_large_number_of_trials_param_in_few_trials():
    # 1000 trials, only 10 have the param, 1 of those has log=True
    trials = [
        FrozenTrial(params={"y": i}, distributions={"y": FloatDistribution(log=False)})
        for i in range(990)
    ] + [
        FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)})
        for i in range(9)
    ] + [
        FrozenTrial(params={"x": 999}, distributions={"x": FloatDistribution(log=True)})
    ]
    codeflash_output = _is_log_scale(trials, "x") # 29.2μs -> 24.8μs (17.5% faster)

def test_large_number_of_trials_param_in_few_trials_none_true():
    # 1000 trials, only 10 have the param, none have log=True
    trials = [
        FrozenTrial(params={"y": i}, distributions={"y": FloatDistribution(log=False)})
        for i in range(990)
    ] + [
        FrozenTrial(params={"x": i}, distributions={"x": FloatDistribution(log=False)})
        for i in range(10)
    ]
    codeflash_output = _is_log_scale(trials, "x") # 28.7μs -> 25.0μs (14.5% 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-_is_log_scale-mhb0vv5v and push.