⚡️ Speed up function _get_pareto_front_plot by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for _get_pareto_front_plot in optuna/visualization/matplotlib/_pareto_front.py

⏱️ Runtime : 255 milliseconds → 232 milliseconds (best of 30 runs)

📝 Explanation and details

The optimization achieves a 9% speedup by replacing Python list comprehensions with NumPy array indexing for extracting coordinate data from trial values.

Key Changes:

• Added import numpy as np
• Replaced list comprehensions like [values[info.axis_order[0]] for _, values in trials] with NumPy array slicing: values_array[:, info.axis_order[0]]
• Applied this pattern to all coordinate extractions in both 2D and 3D plotting functions

Why This is Faster:

1. Vectorized operations: NumPy's array slicing is implemented in C and operates on contiguous memory blocks, avoiding Python's per-element overhead
2. Reduced interpreter overhead: One NumPy operation replaces hundreds/thousands of Python list comprehension iterations
3. Memory efficiency: NumPy arrays store data more compactly than Python lists

Performance Impact:
The optimization shows dramatic improvements for larger datasets:

• Small datasets (1-10 trials): Modest 1-2% improvements
• Medium datasets (100-500 trials): 45-60% speedup
• Large datasets (1000+ trials): Up to 85% speedup

Test Results Analysis:
The annotated tests demonstrate that the optimization is most effective when processing many trials. For example, test_2d_large_number_of_trials shows an 84% improvement (14.5ms → 7.88ms), while smaller test cases show minimal but consistent gains.

This optimization is particularly valuable for Optuna's visualization functionality, as Pareto front plots are commonly used to analyze optimization runs with hundreds or thousands of trials, where the performance gains will be most pronounced.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import builtins
import sys
# Patch points
import types

import matplotlib
import matplotlib.pyplot as plt
# imports
import pytest  # used for our unit tests
from optuna.visualization.matplotlib._pareto_front import \
    _get_pareto_front_plot


# Minimal _ParetoFrontInfo stub
class _ParetoFrontInfo:
    def __init__(
        self,
        n_targets,
        target_names,
        axis_order,
        infeasible_trials_with_values,
        non_best_trials_with_values,
        best_trials_with_values,
    ):
        self.n_targets = n_targets
        self.target_names = target_names
        self.axis_order = axis_order
        self.infeasible_trials_with_values = infeasible_trials_with_values
        self.non_best_trials_with_values = non_best_trials_with_values
        self.best_trials_with_values = best_trials_with_values
from optuna.visualization.matplotlib._pareto_front import \
    _get_pareto_front_plot

# --------------------------
# UNIT TESTS BEGIN HERE
# --------------------------

# --- Basic Test Cases ---

def test_2d_basic_plot_has_labels_and_title():
    # 2D, 1 best, 1 non-best, 1 infeasible
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["x", "y"],
        axis_order=[0, 1],
        infeasible_trials_with_values=[(1, [0.1, 0.2])],
        non_best_trials_with_values=[(2, [0.3, 0.4])],
        best_trials_with_values=[(3, [0.5, 0.6])]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 7.26ms -> 7.24ms (0.322% faster)
    # Check legend exists
    legend = ax.get_legend()

def test_3d_basic_plot_has_labels_and_title():
    # 3D, 2 best, 2 non-best, 2 infeasible
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["x", "y", "z"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=[(1, [0.1, 0.2, 0.3]), (2, [0.4, 0.5, 0.6])],
        non_best_trials_with_values=[(3, [0.7, 0.8, 0.9]), (4, [1.0, 1.1, 1.2])],
        best_trials_with_values=[(5, [1.3, 1.4, 1.5]), (6, [1.6, 1.7, 1.8])]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.5ms -> 10.4ms (1.57% faster)
    # Check legend exists
    legend = ax.get_legend()

def test_2d_only_best_trials():
    # 2D, only best trials
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["foo", "bar"],
        axis_order=[0, 1],
        infeasible_trials_with_values=[],
        non_best_trials_with_values=[],
        best_trials_with_values=[(1, [1.0, 2.0]), (2, [2.0, 3.0])]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 5.20ms -> 5.18ms (0.521% faster)
    # Check legend exists
    legend = ax.get_legend()

def test_3d_only_non_best_trials():
    # 3D, only non-best trials
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["a", "b", "c"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=[],
        non_best_trials_with_values=[(1, [1,2,3]), (2, [4,5,6])],
        best_trials_with_values=[]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 8.31ms -> 8.21ms (1.21% faster)
    # Check legend exists
    legend = ax.get_legend()

# --- Edge Test Cases ---

def test_2d_empty_trials():
    # 2D, all trial lists empty
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["x", "y"],
        axis_order=[0, 1],
        infeasible_trials_with_values=[],
        non_best_trials_with_values=[],
        best_trials_with_values=[]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 3.89ms -> 3.91ms (0.397% slower)

def test_3d_empty_trials():
    # 3D, all trial lists empty
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["x", "y", "z"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=[],
        non_best_trials_with_values=[],
        best_trials_with_values=[]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 6.58ms -> 6.52ms (0.898% faster)

def test_invalid_n_targets_raises():
    # n_targets not 2 or 3 should raise AssertionError
    info = _ParetoFrontInfo(
        n_targets=4,
        target_names=["x", "y", "z", "w"],
        axis_order=[0, 1, 2, 3],
        infeasible_trials_with_values=[],
        non_best_trials_with_values=[],
        best_trials_with_values=[]
    )
    with pytest.raises(AssertionError):
        _get_pareto_front_plot(info) # 1.17μs -> 1.20μs (1.76% slower)

def test_2d_axis_order_swapped():
    # 2D, axis_order swapped
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["first", "second"],
        axis_order=[1, 0],
        infeasible_trials_with_values=[(1, [0.1, 0.2])],
        non_best_trials_with_values=[(2, [0.3, 0.4])],
        best_trials_with_values=[(3, [0.5, 0.6])]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 7.32ms -> 7.24ms (1.12% faster)

def test_3d_axis_order_swapped():
    # 3D, axis_order swapped
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["a", "b", "c"],
        axis_order=[2, 0, 1],
        infeasible_trials_with_values=[(1, [0.1, 0.2, 0.3])],
        non_best_trials_with_values=[(2, [0.4, 0.5, 0.6])],
        best_trials_with_values=[(3, [0.7, 0.8, 0.9])]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.7ms -> 10.5ms (2.00% faster)

def test_2d_infeasible_only():
    # 2D, only infeasible trials
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["x", "y"],
        axis_order=[0, 1],
        infeasible_trials_with_values=[(1, [0.1, 0.2]), (2, [0.3, 0.4])],
        non_best_trials_with_values=[],
        best_trials_with_values=[]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 5.16ms -> 5.12ms (0.766% faster)
    # Legend should exist
    legend = ax.get_legend()

def test_3d_infeasible_only():
    # 3D, only infeasible trials
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["x", "y", "z"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=[(1, [0.1, 0.2, 0.3]), (2, [0.4, 0.5, 0.6])],
        non_best_trials_with_values=[],
        best_trials_with_values=[]
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 8.25ms -> 8.15ms (1.30% faster)
    # Legend should exist
    legend = ax.get_legend()

# --- Large Scale Test Cases ---

def test_2d_large_number_of_trials():
    # 2D, 1000 best, 1000 non-best, 1000 infeasible
    best = [(i, [i*0.1, i*0.2]) for i in range(1000)]
    non_best = [(i, [i*0.3, i*0.4]) for i in range(1000)]
    infeasible = [(i, [i*0.5, i*0.6]) for i in range(1000)]
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["x", "y"],
        axis_order=[0, 1],
        infeasible_trials_with_values=infeasible,
        non_best_trials_with_values=non_best,
        best_trials_with_values=best
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 14.5ms -> 7.88ms (84.1% faster)
    # Check that each scatter plot has 1000 points
    for coll in ax.collections:
        pass

def test_3d_large_number_of_trials():
    # 3D, 1000 best, 1000 non-best, 1000 infeasible
    best = [(i, [i*0.1, i*0.2, i*0.3]) for i in range(1000)]
    non_best = [(i, [i*0.4, i*0.5, i*0.6]) for i in range(1000)]
    infeasible = [(i, [i*0.7, i*0.8, i*0.9]) for i in range(1000)]
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["x", "y", "z"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=infeasible,
        non_best_trials_with_values=non_best,
        best_trials_with_values=best
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 11.5ms -> 11.3ms (2.49% faster)
    # Check that each scatter plot has 1000 points
    for coll in ax.collections:
        pass

def test_2d_performance_large():
    # 2D, test that function runs under 1s for 1000 trials
    import time
    best = [(i, [i*0.1, i*0.2]) for i in range(1000)]
    non_best = [(i, [i*0.3, i*0.4]) for i in range(1000)]
    infeasible = [(i, [i*0.5, i*0.6]) for i in range(1000)]
    info = _ParetoFrontInfo(
        n_targets=2,
        target_names=["x", "y"],
        axis_order=[0, 1],
        infeasible_trials_with_values=infeasible,
        non_best_trials_with_values=non_best,
        best_trials_with_values=best
    )
    t0 = time.time()
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 14.3ms -> 7.74ms (85.3% faster)
    t1 = time.time()

def test_3d_performance_large():
    # 3D, test that function runs under 1s for 1000 trials
    import time
    best = [(i, [i*0.1, i*0.2, i*0.3]) for i in range(1000)]
    non_best = [(i, [i*0.4, i*0.5, i*0.6]) for i in range(1000)]
    infeasible = [(i, [i*0.7, i*0.8, i*0.9]) for i in range(1000)]
    info = _ParetoFrontInfo(
        n_targets=3,
        target_names=["x", "y", "z"],
        axis_order=[0, 1, 2],
        infeasible_trials_with_values=infeasible,
        non_best_trials_with_values=non_best,
        best_trials_with_values=best
    )
    t0 = time.time()
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 11.6ms -> 11.3ms (2.43% faster)
    t1 = time.time()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import matplotlib
import matplotlib.pyplot as plt
# imports
import pytest
from optuna.visualization.matplotlib._pareto_front import \
    _get_pareto_front_plot


# Minimal _ParetoFrontInfo mock for testing
class _ParetoFrontInfo:
    def __init__(
        self,
        n_targets,
        best_trials_with_values,
        non_best_trials_with_values,
        infeasible_trials_with_values,
        target_names,
        axis_order,
    ):
        self.n_targets = n_targets
        self.best_trials_with_values = best_trials_with_values
        self.non_best_trials_with_values = non_best_trials_with_values
        self.infeasible_trials_with_values = infeasible_trials_with_values
        self.target_names = target_names
        self.axis_order = axis_order
from optuna.visualization.matplotlib._pareto_front import \
    _get_pareto_front_plot

# --- Unit tests ---

# 1. Basic Test Cases

def test_2d_basic_single_best_trial():
    # Test with 2D, only one best trial, no non-best or infeasible
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[(1, [0.1, 0.2])],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[],
        target_names=["A", "B"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 5.12ms -> 5.07ms (1.03% faster)
    # Check legend labels
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_2d_basic_best_and_non_best():
    # Test with 2D, one best trial, one non-best, no infeasible
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[(1, [1.0, 2.0])],
        non_best_trials_with_values=[(2, [3.0, 4.0])],
        infeasible_trials_with_values=[],
        target_names=["X", "Y"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 6.17ms -> 6.11ms (1.07% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_2d_basic_all_types():
    # Test with 2D, all trial types present
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[(1, [0.5, 0.5])],
        non_best_trials_with_values=[(2, [0.2, 0.3])],
        infeasible_trials_with_values=[(3, [0.8, 0.9])],
        target_names=["Obj1", "Obj2"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 7.24ms -> 7.13ms (1.56% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_3d_basic_best_trial():
    # Test with 3D, only one best trial, no others
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(1, [0.1, 0.2, 0.3])],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[],
        target_names=["A", "B", "C"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 8.35ms -> 8.23ms (1.44% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_3d_basic_all_types():
    # Test with 3D, all trial types present
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(1, [0.1, 0.2, 0.3])],
        non_best_trials_with_values=[(2, [0.4, 0.5, 0.6])],
        infeasible_trials_with_values=[(3, [0.7, 0.8, 0.9])],
        target_names=["X", "Y", "Z"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.6ms -> 10.4ms (1.51% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

# 2. Edge Test Cases

def test_2d_empty_all_lists():
    # All trial lists are empty
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[],
        target_names=["X", "Y"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 3.92ms -> 3.92ms (0.041% slower)

def test_3d_empty_all_lists():
    # All trial lists are empty in 3D
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[],
        target_names=["X", "Y", "Z"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 6.62ms -> 6.52ms (1.57% faster)

def test_2d_infeasible_only():
    # Only infeasible trials present
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[(1, [0.1, 0.2])],
        target_names=["A", "B"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 5.17ms -> 5.14ms (0.713% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_3d_infeasible_only():
    # Only infeasible trials present in 3D
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[(1, [0.1, 0.2, 0.3])],
        target_names=["A", "B", "C"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 8.26ms -> 8.17ms (1.12% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_2d_axis_order_swapped():
    # axis_order swaps axes
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[(1, [10, 20])],
        non_best_trials_with_values=[(2, [30, 40])],
        infeasible_trials_with_values=[(3, [50, 60])],
        target_names=["First", "Second"],
        axis_order=[1, 0],  # swapped
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 7.34ms -> 7.27ms (1.03% faster)

def test_3d_axis_order_swapped():
    # axis_order swaps axes in 3D
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(1, [1, 2, 3])],
        non_best_trials_with_values=[(2, [4, 5, 6])],
        infeasible_trials_with_values=[(3, [7, 8, 9])],
        target_names=["A", "B", "C"],
        axis_order=[2, 0, 1],  # swapped
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.7ms -> 10.5ms (1.60% faster)

def test_invalid_n_targets_raises():
    # n_targets not 2 or 3 should raise AssertionError
    info = _ParetoFrontInfo(
        n_targets=4,
        best_trials_with_values=[],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[],
        target_names=["A", "B", "C", "D"],
        axis_order=[0, 1, 2, 3],
    )
    with pytest.raises(AssertionError):
        _get_pareto_front_plot(info) # 1.11μs -> 1.09μs (1.56% faster)


def test_3d_infeasible_none():
    # infeasible_trials_with_values is None (should not break)
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(1, [1, 2, 3])],
        non_best_trials_with_values=[(2, [4, 5, 6])],
        infeasible_trials_with_values=None,
        target_names=["A", "B", "C"],
        axis_order=[0, 1, 2],
    )
    # Should not raise
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 9.26ms -> 9.20ms (0.598% faster)
    legend = ax.get_legend()

# 3. Large Scale Test Cases

def test_2d_large_scale():
    # 2D, 500 best, 500 non-best, 500 infeasible
    n = 500
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[(i, [i, i+1]) for i in range(n)],
        non_best_trials_with_values=[(i, [i+2, i+3]) for i in range(n)],
        infeasible_trials_with_values=[(i, [i+4, i+5]) for i in range(n)],
        target_names=["X", "Y"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.8ms -> 7.46ms (45.2% faster)
    # Check that legend exists and all labels are present
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_3d_large_scale():
    # 3D, 300 best, 300 non-best, 300 infeasible
    n = 300
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(i, [i, i+1, i+2]) for i in range(n)],
        non_best_trials_with_values=[(i, [i+3, i+4, i+5]) for i in range(n)],
        infeasible_trials_with_values=[(i, [i+6, i+7, i+8]) for i in range(n)],
        target_names=["X", "Y", "Z"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.7ms -> 10.6ms (1.14% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_2d_large_scale_empty_best():
    # 2D, 900 non-best, 900 infeasible, 0 best
    n = 900
    info = _ParetoFrontInfo(
        n_targets=2,
        best_trials_with_values=[],
        non_best_trials_with_values=[(i, [i, i+1]) for i in range(n)],
        infeasible_trials_with_values=[(i, [i+2, i+3]) for i in range(n)],
        target_names=["A", "B"],
        axis_order=[0, 1],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 10.5ms -> 6.57ms (60.0% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]

def test_3d_large_scale_empty_non_best():
    # 3D, 100 best, 0 non-best, 100 infeasible
    n = 100
    info = _ParetoFrontInfo(
        n_targets=3,
        best_trials_with_values=[(i, [i, i+1, i+2]) for i in range(n)],
        non_best_trials_with_values=[],
        infeasible_trials_with_values=[(i, [i+3, i+4, i+5]) for i in range(n)],
        target_names=["A", "B", "C"],
        axis_order=[0, 1, 2],
    )
    codeflash_output = _get_pareto_front_plot(info); ax = codeflash_output # 9.25ms -> 9.23ms (0.162% faster)
    legend = ax.get_legend()
    labels = [t.get_text() for t in legend.get_texts()]
# 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_pareto_front_plot-mho9wxj5 and push.