⚡️ Speed up function _get_slice_subplot_info by 49%

[codeflash-ai]

📄 49% (0.49x) speedup for _get_slice_subplot_info in optuna/visualization/_slice.py

⏱️ Runtime : 2.41 milliseconds → 1.61 milliseconds (best of 342 runs)

📝 Explanation and details

The optimized code achieves a 49% speedup through three key optimizations:

1. Method Reference Caching

x_append = plot_info.x.append
y_append = plot_info.y.append
# ... etc

This eliminates repeated attribute lookups in the hot loop. Instead of resolving plot_info.x.append on every iteration, the method reference is cached once and reused.

2. Efficient Parameter Lookup

# Original: if param not in t.params: + t.params[param] 
param_value = t.params.get(param, None)
if param_value is None:
    continue

This avoids double dictionary lookups by using .get() with a default value, reducing from two hash table operations to one per trial.

3. Manual Constraint Validation with Early Exit

# Original: all([x <= 0.0 for x in constraints])
if constraints is None or not constraints:
    is_feasible = True
else:
    is_feasible = True
    for x in constraints:
        if x > 0.0:
            is_feasible = False
            break

This eliminates the intermediate list creation from the list comprehension and provides early exit when the first constraint violation is found.

Performance Impact by Test Case:

• Large scale with constraints: Up to 76% faster (constraints with early exit optimization)
• Large scale general: 26-29% faster (method caching benefits)
• Edge cases with None values: 19-517% faster (efficient parameter lookup)
• Small scale: 5-8% slower due to setup overhead, but still functionally identical

The optimizations are most effective for datasets with many trials and constraint validation, which is typical in hyperparameter optimization scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections.abc import Callable
from typing import Any, List, NamedTuple, Optional, Tuple

# imports
import pytest
from optuna.visualization._slice import _get_slice_subplot_info


# Minimal stand-ins for Optuna components for testing purposes
class FrozenTrial:
    def __init__(self, number: int, value: float, params: dict, system_attrs: dict):
        self.number = number
        self.value = value
        self.params = params
        self.system_attrs = system_attrs

# Simulate _CONSTRAINTS_KEY from optuna.samplers._base
_CONSTRAINTS_KEY = "constraints"

# Simulate CategoricalChoiceType from optuna.distributions
CategoricalChoiceType = Any

# Simulate _SliceSubplotInfo from optuna.visualization._slice
class _SliceSubplotInfo(NamedTuple):
    param_name: str
    x: List[Any]
    y: List[float]
    trial_numbers: List[int]
    is_log: bool
    is_numerical: bool
    x_labels: Optional[Tuple[CategoricalChoiceType, ...]]
    constraints: List[bool]
from optuna.visualization._slice import _get_slice_subplot_info

# unit tests

# --- Basic Test Cases ---

def test_basic_single_trial_numerical_param():
    # One trial, param present, numerical value
    trial = FrozenTrial(0, 1.23, {"x": 5}, {})
    codeflash_output = _get_slice_subplot_info([trial], "x", None, False, True, None); info = codeflash_output # 2.99μs -> 3.26μs (8.34% slower)

def test_basic_multiple_trials_mixed_params():
    # Multiple trials, some with param, some without
    trials = [
        FrozenTrial(0, 1.0, {"x": 10}, {}),
        FrozenTrial(1, 2.0, {"y": 20}, {}),
        FrozenTrial(2, 3.0, {"x": 30}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, True, True, None); info = codeflash_output # 3.19μs -> 3.46μs (8.08% slower)

def test_basic_categorical_param_with_labels():
    # Categorical parameter, labels provided
    trials = [
        FrozenTrial(0, 1.0, {"cat": "A"}, {}),
        FrozenTrial(1, 2.0, {"cat": "B"}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "cat", None, False, False, ("A", "B")); info = codeflash_output # 3.04μs -> 3.21μs (5.51% slower)

def test_basic_target_function_override():
    # Use a custom target function
    trials = [
        FrozenTrial(0, 1.0, {"x": 2}, {}),
        FrozenTrial(1, 2.0, {"x": 3}, {}),
    ]
    def target_fn(trial):
        return trial.value * 2
    codeflash_output = _get_slice_subplot_info(trials, "x", target_fn, False, True, None); info = codeflash_output # 3.06μs -> 3.31μs (7.58% slower)

# --- Edge Test Cases ---

def test_edge_no_trials():
    # Empty trials list
    codeflash_output = _get_slice_subplot_info([], "x", None, False, True, None); info = codeflash_output # 1.60μs -> 2.00μs (19.6% slower)

def test_edge_param_not_in_any_trial():
    # Param not present in any trial
    trials = [
        FrozenTrial(0, 1.0, {"y": 2}, {}),
        FrozenTrial(1, 2.0, {"z": 3}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 1.82μs -> 2.23μs (18.1% slower)

def test_edge_param_in_some_trials_only():
    # Param present in some trials, absent in others
    trials = [
        FrozenTrial(0, 1.0, {"x": 5}, {}),
        FrozenTrial(1, 2.0, {"y": 10}, {}),
        FrozenTrial(2, 3.0, {"x": 15}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 3.23μs -> 3.53μs (8.56% slower)

def test_edge_param_value_none():
    # Param present but value is None
    trials = [
        FrozenTrial(0, 1.0, {"x": None}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 2.57μs -> 2.15μs (19.5% faster)

def test_edge_trial_value_none_and_target_fn():
    # Trial value is None, but target function handles it
    trials = [
        FrozenTrial(0, None, {"x": 10}, {}),
    ]
    def target_fn(trial):
        return 42.0  # Arbitrary fallback
    codeflash_output = _get_slice_subplot_info(trials, "x", target_fn, False, True, None); info = codeflash_output # 2.31μs -> 2.50μs (7.60% slower)

def test_edge_constraints_present_and_varied():
    # Constraints present: all <= 0.0 means True, otherwise False
    trials = [
        FrozenTrial(0, 1.0, {"x": 5}, {_CONSTRAINTS_KEY: [-1.0, 0.0]}),  # True
        FrozenTrial(1, 2.0, {"x": 10}, {_CONSTRAINTS_KEY: [0.0, 1.0]}),  # False
        FrozenTrial(2, 3.0, {"x": 15}, {_CONSTRAINTS_KEY: [0.0, 0.0]}),  # True
        FrozenTrial(3, 4.0, {"x": 20}, {_CONSTRAINTS_KEY: [1.0, 2.0]}),  # False
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 5.33μs -> 4.48μs (18.9% faster)

def test_edge_constraints_empty_list():
    # Constraints present but empty list: all([]) is True
    trials = [
        FrozenTrial(0, 1.0, {"x": 5}, {_CONSTRAINTS_KEY: []}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 2.92μs -> 2.75μs (5.88% faster)

def test_edge_constraints_none():
    # system_attrs[_CONSTRAINTS_KEY] is None
    trials = [
        FrozenTrial(0, 1.0, {"x": 5}, {_CONSTRAINTS_KEY: None}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 2.44μs -> 2.68μs (9.17% slower)

def test_edge_param_is_zero_or_negative():
    # Param values are 0 or negative
    trials = [
        FrozenTrial(0, 1.0, {"x": 0}, {}),
        FrozenTrial(1, 2.0, {"x": -5}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 3.06μs -> 3.21μs (4.92% slower)

def test_edge_param_is_float_and_int_and_str():
    # Param is float, int, and str across trials
    trials = [
        FrozenTrial(0, 1.0, {"x": 1.5}, {}),
        FrozenTrial(1, 2.0, {"x": 2}, {}),
        FrozenTrial(2, 3.0, {"x": "foo"}, {}),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 3.25μs -> 3.41μs (4.43% slower)

# --- Large Scale Test Cases ---

def test_large_scale_many_trials_numerical():
    # 1000 trials, all with param "x", numerical values
    trials = [FrozenTrial(i, float(i), {"x": i}, {}) for i in range(1000)]
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 257μs -> 203μs (26.4% faster)

def test_large_scale_some_trials_missing_param():
    # 1000 trials, half with param "x", half without
    trials = []
    for i in range(1000):
        if i % 2 == 0:
            trials.append(FrozenTrial(i, float(i), {"x": i}, {}))
        else:
            trials.append(FrozenTrial(i, float(i), {"y": i}, {}))
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 143μs -> 126μs (12.8% faster)

def test_large_scale_constraints_varied():
    # 1000 trials, constraints alternate between True/False
    trials = []
    for i in range(1000):
        if i % 2 == 0:
            constraints = [-1.0, 0.0]  # True
        else:
            constraints = [1.0, 0.0]   # False
        trials.append(FrozenTrial(i, float(i), {"x": i}, {_CONSTRAINTS_KEY: constraints}))
    codeflash_output = _get_slice_subplot_info(trials, "x", None, False, True, None); info = codeflash_output # 431μs -> 245μs (76.2% faster)

def test_large_scale_categorical_param():
    # 1000 trials, categorical param cycling through 10 categories
    categories = tuple(f"cat{i}" for i in range(10))
    trials = []
    for i in range(1000):
        cat = categories[i % 10]
        trials.append(FrozenTrial(i, float(i), {"c": cat}, {}))
    codeflash_output = _get_slice_subplot_info(trials, "c", None, False, False, categories); info = codeflash_output # 255μs -> 202μs (26.3% faster)

def test_large_scale_target_function():
    # 1000 trials, custom target function
    trials = [FrozenTrial(i, float(i), {"x": i}, {}) for i in range(1000)]
    def target_fn(trial):
        return trial.value + 1000
    codeflash_output = _get_slice_subplot_info(trials, "x", target_fn, False, True, None); info = codeflash_output # 238μs -> 183μs (29.7% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Any, Callable, Dict, List, Optional, Tuple

# imports
import pytest
from optuna.visualization._slice import _get_slice_subplot_info


# Minimal stand-ins for optuna classes/types for testing
class FrozenTrial:
    def __init__(
        self,
        number: int,
        params: Dict[str, Any],
        value: Optional[float] = None,
        system_attrs: Optional[Dict[str, Any]] = None,
    ):
        self.number = number
        self.params = params
        self.value = value
        self.system_attrs = system_attrs or {}

# Stand-in for CategoricalChoiceType (accepts anything)
CategoricalChoiceType = Any

# Stand-in for _CONSTRAINTS_KEY
_CONSTRAINTS_KEY = "constraints"

# Stand-in for _SliceSubplotInfo
class _SliceSubplotInfo:
    def __init__(
        self,
        param_name: str,
        x: List[Any],
        y: List[float],
        trial_numbers: List[int],
        is_log: bool,
        is_numerical: bool,
        x_labels: Optional[Tuple[CategoricalChoiceType, ...]],
        constraints: List[bool],
    ):
        self.param_name = param_name
        self.x = x
        self.y = y
        self.trial_numbers = trial_numbers
        self.is_log = is_log
        self.is_numerical = is_numerical
        self.x_labels = x_labels
        self.constraints = constraints

    def __eq__(self, other):
        if not isinstance(other, _SliceSubplotInfo):
            return False
        return (
            self.param_name == other.param_name
            and self.x == other.x
            and self.y == other.y
            and self.trial_numbers == other.trial_numbers
            and self.is_log == other.is_log
            and self.is_numerical == other.is_numerical
            and self.x_labels == other.x_labels
            and self.constraints == other.constraints
        )
from optuna.visualization._slice import _get_slice_subplot_info

# unit tests

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

def test_basic_single_trial_numerical_param():
    # One trial, param present, no constraints, default target (value)
    trial = FrozenTrial(number=1, params={"a": 5}, value=42.0)
    codeflash_output = _get_slice_subplot_info([trial], "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 3.55μs -> 3.76μs (5.50% slower)

def test_basic_multiple_trials_mixed_params():
    # Multiple trials, some with param, some without
    trials = [
        FrozenTrial(number=0, params={"a": 1, "b": 2}, value=10.0),
        FrozenTrial(number=1, params={"b": 3}, value=20.0),
        FrozenTrial(number=2, params={"a": 4}, value=30.0),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=True, numerical=False, x_labels=("one", "four")); result = codeflash_output # 3.66μs -> 3.92μs (6.56% slower)

def test_basic_with_custom_target():
    # Use a custom target function
    trials = [
        FrozenTrial(number=0, params={"a": 2}, value=3.0),
        FrozenTrial(number=1, params={"a": 5}, value=7.0),
    ]
    def target_fn(trial):
        return trial.value * 2
    codeflash_output = _get_slice_subplot_info(trials, "a", target_fn, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 3.32μs -> 3.57μs (6.81% slower)

def test_basic_categorical_param():
    # Param is categorical
    trials = [
        FrozenTrial(number=0, params={"cat": "foo"}, value=1.1),
        FrozenTrial(number=1, params={"cat": "bar"}, value=2.2),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "cat", None, log_scale=False, numerical=False, x_labels=("foo", "bar")); result = codeflash_output # 3.39μs -> 3.66μs (7.25% slower)

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

def test_edge_no_trials():
    # No trials at all
    codeflash_output = _get_slice_subplot_info([], "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 1.88μs -> 2.23μs (15.8% slower)

def test_edge_param_never_present():
    # Trials never have the requested param
    trials = [
        FrozenTrial(number=0, params={"b": 1}, value=1.0),
        FrozenTrial(number=1, params={"c": 2}, value=2.0),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 2.04μs -> 2.42μs (15.8% slower)


def test_edge_constraints_key_various_cases():
    # Constraints absent, present and all <=0, present and some >0
    trials = [
        FrozenTrial(number=0, params={"a": 1}, value=1.0),  # No constraints
        FrozenTrial(number=1, params={"a": 2}, value=2.0, system_attrs={_CONSTRAINTS_KEY: [-1.0, 0.0]}),  # valid
        FrozenTrial(number=2, params={"a": 3}, value=3.0, system_attrs={_CONSTRAINTS_KEY: [0.0, 1.0]}),   # invalid
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 6.06μs -> 5.23μs (16.0% faster)

def test_edge_empty_params_dict():
    # Trial has empty params dict
    trials = [
        FrozenTrial(number=0, params={}, value=1.0),
        FrozenTrial(number=1, params={"a": 2}, value=2.0),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 3.30μs -> 3.52μs (6.16% slower)

def test_edge_param_with_none_value():
    # Param value is None (should be included as x=None)
    trials = [
        FrozenTrial(number=0, params={"a": None}, value=7.0),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 2.86μs -> 2.56μs (11.7% faster)

def test_edge_param_with_various_types():
    # Param values are int, float, str, bool
    trials = [
        FrozenTrial(number=0, params={"a": 1}, value=1.0),
        FrozenTrial(number=1, params={"a": 2.5}, value=2.0),
        FrozenTrial(number=2, params={"a": "foo"}, value=3.0),
        FrozenTrial(number=3, params={"a": True}, value=4.0),
    ]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 4.02μs -> 4.14μs (2.92% slower)

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

def test_large_many_trials():
    # 1000 trials, all with param, values from 0 to 999
    n = 1000
    trials = [FrozenTrial(number=i, params={"a": i}, value=float(i)) for i in range(n)]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=True, numerical=True, x_labels=None); result = codeflash_output # 261μs -> 205μs (27.4% faster)

def test_large_mixed_constraints():
    # 500 trials with constraints all <=0, 500 with some >0
    n = 500
    trials = []
    for i in range(n):
        trials.append(FrozenTrial(number=i, params={"a": i}, value=float(i), system_attrs={_CONSTRAINTS_KEY: [-1.0, 0.0]}))
    for i in range(n, 2*n):
        trials.append(FrozenTrial(number=i, params={"a": i}, value=float(i), system_attrs={_CONSTRAINTS_KEY: [0.0, 1.0]}))
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 435μs -> 260μs (67.4% faster)

def test_large_sparse_param():
    # 1000 trials, only every 10th trial has the param
    n = 1000
    trials = []
    for i in range(n):
        if i % 10 == 0:
            trials.append(FrozenTrial(number=i, params={"a": i}, value=float(i)))
        else:
            trials.append(FrozenTrial(number=i, params={"b": i}, value=float(i)))
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 50.2μs -> 66.0μs (23.9% slower)

def test_large_all_params_none():
    # 1000 trials, all param values are None
    n = 1000
    trials = [FrozenTrial(number=i, params={"a": None}, value=float(i)) for i in range(n)]
    codeflash_output = _get_slice_subplot_info(trials, "a", None, log_scale=False, numerical=True, x_labels=None); result = codeflash_output # 262μs -> 42.5μs (517% 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-_get_slice_subplot_info-mhjo7f2b and push.