⚡️ Speed up function periodogram by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for periodogram in quantecon/_estspec.py

⏱️ Runtime : 1.62 milliseconds → 1.53 milliseconds (best of 290 runs)

📝 Explanation and details

The optimized code achieves a 6% speedup through three key performance improvements:

1. Real FFT optimization in periodogram

• Replaced fft(x) with np.fft.rfft(x) for real input data
• rfft computes only frequencies in [0, π], halving both computation and memory usage
• Eliminates the need for slicing [:int(n/2)+1] since rfft naturally produces the correct output size
• Major impact: Line profiler shows FFT computation drops from 59.3% to 59.9% of total time, but with reduced absolute time

2. Direct window function lookup in smooth

• Replaced dictionary lookup (windows[window](window_len)) with direct if/elif chain
• Eliminates dictionary creation overhead (was 2.4% of smooth runtime) and function pointer dereferencing
• Performance gain: Most evident in test cases with windowing - up to 8.55% faster for individual windowed calls

3. Efficient array concatenation

• Replaced np.concatenate((xb[::-1], x, xt[::-1])) with pre-allocated np.empty buffer and direct slice assignments
• Avoids creating temporary arrays and reduces memory allocations
• Impact: Concatenation overhead drops from 12.7% to combined 11.4% (2.8% + 2.2% + 2.8% + 2.8%) in line profiler

Test case benefits:

• Large arrays see the biggest gains (15-20% faster) due to FFT and memory allocation improvements
• Windowed operations benefit from direct lookup (4-8% faster)
• Small arrays show modest improvements (1-5%) as overhead reduction is less significant

The optimizations are particularly effective for the common use case of large signal processing tasks with optional smoothing windows.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from numpy.fft import fft
from quantecon._estspec import periodogram

# unit tests

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

def test_periodogram_basic_constant():
    # Test with a constant input array
    x = np.ones(8)
    w, I_w = periodogram(x) # 39.2μs -> 38.8μs (1.11% faster)

def test_periodogram_basic_sine():
    # Test with a pure sine wave
    n = 16
    t = np.arange(n)
    freq = 2
    x = np.sin(2 * np.pi * freq * t / n)
    w, I_w = periodogram(x) # 24.4μs -> 23.6μs (3.76% faster)
    # The periodogram should peak at the correct frequency bin
    peak_bin = np.argmax(I_w)
    expected_bin = freq

def test_periodogram_basic_window_flat():
    # Test with smoothing window 'flat'
    x = np.random.rand(20)
    w, I_w = periodogram(x, window='flat', window_len=5) # 52.2μs -> 49.3μs (5.95% faster)
    # Output length should be len(x)//2+1 minus window_len//2*2 (due to smoothing)
    expected_len = len(x)//2 + 1 - (5//2)*2

def test_periodogram_basic_window_hanning():
    # Test with smoothing window 'hanning'
    x = np.random.rand(20)
    w, I_w = periodogram(x, window='hanning', window_len=7) # 54.8μs -> 52.4μs (4.46% faster)
    expected_len = len(x)//2 + 1 - (7//2)*2

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

def test_periodogram_edge_empty():
    # Test with empty input
    x = np.array([])
    with pytest.raises(ValueError):
        periodogram(x) # 3.25μs -> 3.67μs (11.4% slower)

def test_periodogram_edge_single_element():
    # Test with single element input
    x = np.array([42.0])
    w, I_w = periodogram(x) # 29.9μs -> 29.7μs (0.785% faster)

def test_periodogram_edge_two_elements():
    # Test with two elements
    x = np.array([1.0, -1.0])
    w, I_w = periodogram(x) # 25.5μs -> 24.8μs (2.73% faster)

def test_periodogram_edge_non_power_of_two():
    # Test with a length that is not a power of two
    x = np.random.rand(15)
    w, I_w = periodogram(x) # 26.4μs -> 25.4μs (4.07% faster)

def test_periodogram_edge_negative_values():
    # Test with negative values
    x = -np.ones(8)
    w, I_w = periodogram(x) # 22.8μs -> 22.7μs (0.705% faster)

def test_periodogram_edge_window_len_too_large():
    # Test window_len larger than input
    x = np.random.rand(5)
    with pytest.raises(ValueError):
        periodogram(x, window='hanning', window_len=11) # 26.6μs -> 26.4μs (0.983% faster)

def test_periodogram_edge_window_len_even():
    # Test with even window_len (should auto-increment to odd)
    x = np.random.rand(20)
    w, I_w = periodogram(x, window='hanning', window_len=6) # 61.4μs -> 59.3μs (3.47% faster)
    # Output length should reflect incremented window_len
    expected_len = len(x)//2 + 1 - (7//2)*2  # window_len auto-incremented to 7

def test_periodogram_edge_window_len_too_small():
    # Test window_len less than 3
    x = np.random.rand(10)
    with pytest.raises(ValueError):
        periodogram(x, window='hanning', window_len=2) # 27.0μs -> 27.3μs (0.938% slower)

def test_periodogram_edge_unrecognized_window():
    # Test with unrecognized window type (should default to hanning)
    x = np.random.rand(20)
    w, I_w = periodogram(x, window='unknown', window_len=7) # 58.0μs -> 56.6μs (2.53% faster)
    # Output should still be correct length
    expected_len = len(x)//2 + 1 - (7//2)*2

def test_periodogram_edge_nan_inf():
    # Test with NaN and Inf values
    x = np.array([1.0, np.nan, 2.0, np.inf, 3.0, -np.inf])
    w, I_w = periodogram(x) # 25.1μs -> 35.6μs (29.6% slower)

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

def test_periodogram_large_random():
    # Test with a large random array
    x = np.random.rand(1000)
    w, I_w = periodogram(x) # 39.9μs -> 34.6μs (15.5% faster)

def test_periodogram_large_sine_peak():
    # Test with a large sine wave, peak should be at correct frequency
    n = 1000
    freq = 50
    t = np.arange(n)
    x = np.sin(2 * np.pi * freq * t / n)
    w, I_w = periodogram(x) # 34.7μs -> 29.3μs (18.6% faster)
    peak_bin = np.argmax(I_w)

def test_periodogram_large_window_flat():
    # Test with large input and smoothing
    x = np.random.rand(1000)
    w, I_w = periodogram(x, window='flat', window_len=21) # 69.2μs -> 61.2μs (13.1% faster)
    expected_len = len(x)//2 + 1 - (21//2)*2

def test_periodogram_large_window_hamming():
    # Test with large input and hamming window smoothing
    x = np.random.rand(1000)
    w, I_w = periodogram(x, window='hamming', window_len=15) # 74.2μs -> 65.9μs (12.7% faster)
    expected_len = len(x)//2 + 1 - (15//2)*2
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import math

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from numpy.fft import fft
from quantecon._estspec import periodogram

# =======================
# Unit Tests for periodogram
# =======================

# Helper to compare arrays with relative tolerance, since periodogram is numerical
def arrays_close(a, b, rtol=1e-7, atol=1e-10):
    for i in range(len(a)):
        pass

# ---- 1. BASIC TEST CASES ----

def test_periodogram_constant_signal():
    # Constant signal should have all energy at zero frequency
    x = np.ones(16)
    w, I_w = periodogram(x) # 31.5μs -> 31.5μs (0.105% faster)
    for i in range(1, len(I_w)):
        pass

def test_periodogram_single_sine():
    # Sine wave at a single frequency should have a peak at the corresponding frequency
    n = 32
    freq = 3  # cycles per n samples
    t = np.arange(n)
    x = np.sin(2 * np.pi * freq * t / n)
    w, I_w = periodogram(x) # 23.6μs -> 22.6μs (4.55% faster)
    # Find peak
    peak_idx = np.argmax(I_w)
    # The peak should be much larger than the mean of the rest
    mean_rest = (np.sum(I_w) - I_w[peak_idx]) / (len(I_w) - 1)

def test_periodogram_real_output_length():
    # For real input, output length should be n//2+1
    for n in [1, 2, 3, 8, 9, 16, 17, 100]:
        x = np.random.randn(n)
        w, I_w = periodogram(x) # 91.4μs -> 86.3μs (5.84% faster)

def test_periodogram_window_flat():
    # Test with window smoothing, flat window (moving average)
    n = 32
    x = np.random.randn(n)
    w1, I1 = periodogram(x) # 24.3μs -> 23.4μs (3.90% faster)
    w2, I2 = periodogram(x, window="flat", window_len=5) # 36.4μs -> 34.2μs (6.31% faster)

def test_periodogram_window_hanning():
    # Test with hanning window smoothing
    n = 32
    x = np.random.randn(n)
    w1, I1 = periodogram(x) # 24.2μs -> 23.4μs (3.68% faster)
    w2, I2 = periodogram(x, window="hanning", window_len=5) # 39.0μs -> 37.1μs (5.08% faster)

def test_periodogram_window_invalid_type_defaults():
    # Invalid window type should default to hanning
    n = 32
    x = np.random.randn(n)
    w1, I1 = periodogram(x, window="notawindow", window_len=5) # 51.8μs -> 49.5μs (4.70% faster)
    w2, I2 = periodogram(x, window="hanning", window_len=5) # 26.4μs -> 24.7μs (6.76% faster)
    arrays_close(I1, I2)

def test_periodogram_window_even_length():
    # Even window_len should be promoted to next odd value
    n = 32
    x = np.random.randn(n)
    w1, I1 = periodogram(x, window="hanning", window_len=6) # 49.2μs -> 47.7μs (3.13% faster)
    w2, I2 = periodogram(x, window="hanning", window_len=7) # 25.6μs -> 23.6μs (8.55% faster)
    arrays_close(I1, I2)

# ---- 2. EDGE TEST CASES ----

def test_periodogram_empty_input():
    # Should raise an error for empty input
    with pytest.raises(ValueError):
        periodogram(np.array([])) # 3.10μs -> 3.45μs (10.2% slower)

def test_periodogram_single_element():
    # Should work for single element
    x = np.array([42.])
    w, I_w = periodogram(x) # 27.8μs -> 27.8μs (0.104% slower)

def test_periodogram_two_elements():
    # Should work for two elements
    x = np.array([1., -1.])
    w, I_w = periodogram(x) # 23.7μs -> 23.0μs (3.04% faster)
    # The periodogram should be nonnegative
    for val in I_w:
        pass

def test_periodogram_window_len_too_large():
    # Smoothing window longer than input should raise ValueError
    x = np.random.randn(5)
    with pytest.raises(ValueError):
        periodogram(x, window="flat", window_len=11) # 27.3μs -> 26.5μs (3.19% faster)

def test_periodogram_window_len_too_small():
    # Window length less than 3 should raise ValueError
    x = np.random.randn(10)
    with pytest.raises(ValueError):
        periodogram(x, window="hanning", window_len=2) # 27.4μs -> 26.9μs (1.93% faster)

def test_periodogram_negative_values():
    # Should work for negative values
    x = -np.abs(np.random.randn(16))
    w, I_w = periodogram(x) # 23.5μs -> 24.1μs (2.69% slower)
    # All periodogram values should be nonnegative
    for val in I_w:
        pass

def test_periodogram_nan_input():
    # Should propagate NaN in input to output
    x = np.random.randn(16)
    x[5] = float('nan')
    w, I_w = periodogram(x) # 24.8μs -> 24.3μs (2.12% faster)

def test_periodogram_inf_input():
    # Should propagate inf in input to output
    x = np.random.randn(16)
    x[3] = float('inf')
    w, I_w = periodogram(x) # 35.7μs -> 23.8μs (49.9% faster)

def test_periodogram_all_zeros():
    # All zeros should produce all zeros in output
    x = np.zeros(20)
    w, I_w = periodogram(x) # 25.9μs -> 24.1μs (7.30% faster)
    for val in I_w:
        pass

def test_periodogram_nonfloat_input():
    # Should work with integer input
    x = np.arange(16)
    w, I_w = periodogram(x) # 24.1μs -> 24.1μs (0.141% slower)
    for val in I_w:
        pass

def test_periodogram_minimal_window_len():
    # Should work with minimal valid window length
    x = np.random.randn(7)
    w, I_w = periodogram(x, window="flat", window_len=3) # 49.5μs -> 48.5μs (2.03% faster)

# ---- 3. LARGE SCALE TEST CASES ----

def test_periodogram_large_random_signal():
    # Should work efficiently for large n (n=1000)
    n = 1000
    x = np.random.randn(n)
    w, I_w = periodogram(x) # 40.7μs -> 34.7μs (17.2% faster)
    for val in I_w:
        pass

def test_periodogram_large_sine_superposition():
    # Superposition of multiple sines, check for multiple peaks
    n = 1000
    t = np.arange(n)
    x = np.sin(2 * np.pi * 5 * t / n) + 0.5 * np.sin(2 * np.pi * 20 * t / n)
    w, I_w = periodogram(x) # 34.3μs -> 28.6μs (20.1% faster)
    # Find two largest peaks
    idxs = np.argsort(I_w)[-2:]

def test_periodogram_large_with_window():
    # Large signal with window smoothing
    n = 1000
    x = np.random.randn(n)
    w1, I1 = periodogram(x) # 39.5μs -> 33.5μs (18.0% faster)
    w2, I2 = periodogram(x, window="hamming", window_len=51) # 62.2μs -> 55.4μs (12.3% faster)

def test_periodogram_large_all_zeros():
    # Large all-zero input
    n = 1000
    x = np.zeros(n)
    w, I_w = periodogram(x) # 36.3μs -> 31.0μs (16.9% faster)
    for val in I_w:
        pass

To edit these changes git checkout codeflash/optimize-periodogram-mgh0yu30 and push.