oscura 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

oscura/analyzers/statistics/correlation.py

@@ -25,12 +25,68 @@ from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 
+from oscura.core.numba_backend import njit, prange
 from oscura.core.types import WaveformTrace
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
 
+@njit(parallel=True, cache=True)  # type: ignore[untyped-decorator]  # Numba JIT decorator
+def _autocorr_direct_numba(
+    data: np.ndarray,  # type: ignore[type-arg]
+    max_lag: int,
+) -> np.ndarray:  # type: ignore[type-arg]
+    """Compute autocorrelation using direct method with Numba JIT compilation.
+
+    Alternative implementation using Numba for autocorrelation computation.
+    Uses parallel execution across lags for potential speedup on multi-core systems.
+
+    **Note**: In practice, numpy.correlate is faster for most cases due to highly
+    optimized BLAS/LAPACK routines. This function is provided for educational
+    purposes and specific use cases where custom computation logic is needed.
+
+    Args:
+        data: Mean-centered input signal data (1D array).
+        max_lag: Maximum lag to compute (inclusive).
+
+    Returns:
+        Autocorrelation values from lag 0 to max_lag (unnormalized).
+
+    Performance characteristics:
+        - First call: ~100-200ms compilation overhead (cached for subsequent calls)
+        - Typical performance: ~2ms for n=100, ~10ms for n=256 (compiled)
+        - NumPy's correlate: ~0.02ms for n=100 (100x faster due to BLAS)
+        - Parallel execution: Benefits from multi-core for large max_lag
+
+    Example:
+        >>> import numpy as np
+        >>> data = np.random.randn(128) - np.mean(np.random.randn(128))
+        >>> acf = _autocorr_direct_numba(data, max_lag=64)
+        >>> print(acf.shape)  # (65,)
+
+    Notes:
+        - Input data should be mean-centered before calling this function
+        - Result is NOT normalized; caller must normalize if needed
+        - NumPy's correlate is recommended for production use (faster)
+        - Thread safety: Numba releases GIL, safe for parallel execution
+
+    References:
+        Box, G. E. P. & Jenkins, G. M. (1976). Time Series Analysis
+    """
+    n = len(data)
+    acf = np.zeros(max_lag + 1, dtype=np.float64)
+
+    # Compute autocorrelation for each lag in parallel
+    for lag in prange(max_lag + 1):
+        sum_val = 0.0
+        for i in range(n - lag):
+            sum_val += data[i] * data[i + lag]
+        acf[lag] = sum_val
+
+    return acf
+
+
 @dataclass
 class CrossCorrelationResult:
     """Result of cross-correlation analysis.
@@ -66,6 +122,10 @@ def autocorrelation(
     Measures self-similarity of a signal at different time lags.
     Useful for detecting periodicities and characteristic time scales.
 
+    This function automatically selects the optimal computation method:
+    - Small signals (n < 256): Direct method using numpy.correlate (optimized BLAS)
+    - Large signals (n >= 256): FFT-based method (O(n log n) complexity)
+
     Args:
         trace: Input trace or numpy array.
         max_lag: Maximum lag to compute (samples). If None, uses n // 2.
@@ -80,6 +140,11 @@ def autocorrelation(
     Raises:
         ValueError: If sample_rate is not provided when trace is array.
 
+    Performance:
+        - Small signals (n<256): ~0.02-0.05ms using numpy.correlate
+        - Large signals (n>=256): ~0.1-1ms using FFT method
+        - Both methods use highly optimized numerical libraries
+
     Example:
         >>> lag_times, acf = autocorrelation(trace, max_lag=1000)
         >>> # Find first zero crossing for decorrelation time
@@ -108,15 +173,19 @@ def autocorrelation(
     # Remove mean for proper correlation
     data_centered = data - np.mean(data)
 
-    # Compute autocorrelation via FFT (faster for large n)
-    if n > 256:
+    # Compute autocorrelation using optimal method based on signal size:
+    # - n >= 256: FFT-based method (O(n log n) complexity, fastest for large signals)
+    # - n < 256: Direct method using numpy.correlate (highly optimized BLAS)
+    # Note: Numba implementation available (_autocorr_direct_numba) but numpy.correlate
+    #       is faster due to optimized BLAS/LAPACK routines
+    if n >= 256:
         # Zero-pad for full correlation
         nfft = int(2 ** np.ceil(np.log2(2 * n)))
         fft_data = np.fft.rfft(data_centered, n=nfft)
         acf_full = np.fft.irfft(fft_data * np.conj(fft_data), n=nfft)
         acf = acf_full[: max_lag + 1]
     else:
-        # Direct computation for small n
+        # Direct computation for small signals (numpy.correlate uses optimized BLAS)
         acf = np.correlate(data_centered, data_centered, mode="full")
         acf = acf[n - 1 : n + max_lag]
 
@@ -131,6 +200,71 @@ def autocorrelation(
     return lag_times, acf.astype(np.float64)
 
 
+def _extract_correlation_data(
+    trace1: WaveformTrace | NDArray[np.floating[Any]],
+    trace2: WaveformTrace | NDArray[np.floating[Any]],
+    sample_rate: float | None,
+) -> tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]], float]:
+    """Extract data arrays and sample rate from traces.
+
+    Args:
+        trace1: First input trace or array
+        trace2: Second input trace or array
+        sample_rate: Sample rate if traces are arrays
+
+    Returns:
+        Tuple of (data1, data2, sample_rate)
+
+    Raises:
+        ValueError: If sample_rate needed but not provided
+    """
+    if isinstance(trace1, WaveformTrace):
+        data1 = trace1.data
+        fs = trace1.metadata.sample_rate
+    else:
+        data1 = trace1
+        if sample_rate is None:
+            raise ValueError("sample_rate required when traces are arrays")
+        fs = sample_rate
+
+    if isinstance(trace2, WaveformTrace):
+        data2 = trace2.data
+        if not isinstance(trace1, WaveformTrace):
+            fs = trace2.metadata.sample_rate
+    else:
+        data2 = trace2
+
+    return data1, data2, fs
+
+
+def _compute_normalized_xcorr(
+    data1_centered: NDArray[np.floating[Any]],
+    data2_centered: NDArray[np.floating[Any]],
+    xcorr: NDArray[np.floating[Any]],
+) -> NDArray[np.floating[Any]]:
+    """Normalize cross-correlation to [-1, 1].
+
+    Args:
+        data1_centered: First centered data array
+        data2_centered: Second centered data array
+        xcorr: Raw cross-correlation
+
+    Returns:
+        Normalized cross-correlation
+    """
+    norm1 = np.sqrt(np.sum(data1_centered**2))
+    norm2 = np.sqrt(np.sum(data2_centered**2))
+    if norm1 > 0 and norm2 > 0:
+        # Division returns proper NDArray type
+        if TYPE_CHECKING:
+            from typing import cast
+
+            return cast("NDArray[np.floating[Any]]", xcorr / (norm1 * norm2))
+        else:
+            return xcorr / (norm1 * norm2)  # type: ignore[return-value]
+    return xcorr
+
+
 def cross_correlation(
     trace1: WaveformTrace | NDArray[np.floating[Any]],
     trace2: WaveformTrace | NDArray[np.floating[Any]],
@@ -165,71 +299,37 @@ def cross_correlation(
     References:
         Oppenheim, A. V. & Schafer, R. W. (2009). Discrete-Time Signal Processing
     """
-    if isinstance(trace1, WaveformTrace):
-        data1 = trace1.data
-        fs = trace1.metadata.sample_rate
-    else:
-        data1 = trace1
-        if sample_rate is None:
-            raise ValueError("sample_rate required when traces are arrays")
-        fs = sample_rate
+    data1, data2, fs = _extract_correlation_data(trace1, trace2, sample_rate)
+    n1 = len(data1)
+    max_lag = min(len(data1), len(data2)) // 2 if max_lag is None else max_lag
 
-    if isinstance(trace2, WaveformTrace):
-        data2 = trace2.data
-        # Use trace2 sample rate if available and trace1 wasn't a WaveformTrace
-        if not isinstance(trace1, WaveformTrace):
-            fs = trace2.metadata.sample_rate
-    else:
-        data2 = trace2
-
-    n1, n2 = len(data1), len(data2)
-
-    if max_lag is None:
-        max_lag = min(n1, n2) // 2
-
-    # Center the data
+    # Center and compute correlation
     data1_centered = data1 - np.mean(data1)
     data2_centered = data2 - np.mean(data2)
-
-    # Full cross-correlation
-    # Note: np.correlate(a, b) computes sum(a[n+k] * conj(b[k]))
-    # For cross-correlation where we want to detect b delayed relative to a,
-    # we need correlate(b, a) so positive lag means b is delayed
     xcorr_full = np.correlate(data2_centered, data1_centered, mode="full")
 
-    # Extract relevant portion around zero lag
-    # Full correlation has length n1 + n2 - 1, with zero lag at index n1 - 1
-    # (since we swapped the order above)
+    # Extract relevant portion
     zero_lag_idx = n1 - 1
     start_idx = max(0, zero_lag_idx - max_lag)
     end_idx = min(len(xcorr_full), zero_lag_idx + max_lag + 1)
     xcorr = xcorr_full[start_idx:end_idx]
-
-    # Create lag array
     lags = np.arange(start_idx - zero_lag_idx, end_idx - zero_lag_idx)
 
-    # Normalize
+    # Normalize if requested
     if normalized:
-        norm1 = np.sqrt(np.sum(data1_centered**2))
-        norm2 = np.sqrt(np.sum(data2_centered**2))
-        if norm1 > 0 and norm2 > 0:
-            xcorr = xcorr / (norm1 * norm2)
+        xcorr = _compute_normalized_xcorr(data1_centered, data2_centered, xcorr)
 
     # Find peak
     peak_local_idx = np.argmax(np.abs(xcorr))
     peak_lag = int(lags[peak_local_idx])
     peak_coefficient = float(xcorr[peak_local_idx])
 
-    # Time values
-    lag_times = lags / fs
-    peak_lag_time = peak_lag / fs
-
     return CrossCorrelationResult(
         correlation=xcorr.astype(np.float64),
         lags=lags,
-        lag_times=lag_times.astype(np.float64),
+        lag_times=(lags / fs).astype(np.float64),
         peak_lag=peak_lag,
-        peak_lag_time=peak_lag_time,
+        peak_lag_time=peak_lag / fs,
         peak_coefficient=peak_coefficient,
         sample_rate=fs,
     )
@@ -267,6 +367,92 @@ def correlation_coefficient(
     return float(np.corrcoef(data1, data2)[0, 1])
 
 
+def _extract_periodicity_data(
+    trace: WaveformTrace | NDArray[np.floating[Any]], sample_rate: float | None
+) -> tuple[NDArray[np.floating[Any]], float]:
+    """Extract data and sample rate from trace.
+
+    Args:
+        trace: Input trace or array.
+        sample_rate: Sample rate if array.
+
+    Returns:
+        Tuple of (data, sample_rate).
+
+    Raises:
+        ValueError: If sample_rate not provided for array.
+    """
+    if isinstance(trace, WaveformTrace):
+        return trace.data, trace.metadata.sample_rate
+    else:
+        if sample_rate is None:
+            raise ValueError("sample_rate required when trace is array")
+        return trace, sample_rate
+
+
+def _find_primary_peak(acf: NDArray[np.float64], min_period_samples: int) -> tuple[int, float]:
+    """Find primary peak in autocorrelation function.
+
+    Args:
+        acf: Autocorrelation function.
+        min_period_samples: Minimum period to consider.
+
+    Returns:
+        Tuple of (period_samples, strength).
+    """
+    acf_search = acf[min_period_samples:]
+
+    if len(acf_search) < 3:
+        return -1, np.nan
+
+    local_max = (acf_search[1:-1] > acf_search[:-2]) & (acf_search[1:-1] > acf_search[2:])
+    max_indices = np.where(local_max)[0] + 1
+
+    if len(max_indices) == 0:
+        primary_idx = int(np.argmax(acf_search)) + min_period_samples
+    else:
+        peak_values = acf_search[max_indices]
+        best_peak_idx = int(np.argmax(peak_values))
+        primary_idx = int(max_indices[best_peak_idx]) + min_period_samples
+
+    return primary_idx, float(acf[primary_idx])
+
+
+def _find_harmonics(acf: NDArray[np.float64], period_samples: int) -> list[dict[str, int | float]]:
+    """Find harmonic peaks at multiples of primary period.
+
+    Args:
+        acf: Autocorrelation function.
+        period_samples: Primary period in samples.
+
+    Returns:
+        List of harmonic dictionaries.
+    """
+    harmonics: list[dict[str, int | float]] = []
+
+    for h in range(2, 6):
+        harmonic_lag = h * period_samples
+        if harmonic_lag >= len(acf):
+            break
+
+        search_range = max(1, period_samples // 4)
+        start = int(max(0, harmonic_lag - search_range))
+        end = int(min(len(acf), harmonic_lag + search_range))
+        local_max_idx = int(start + int(np.argmax(acf[start:end])))
+        harmonic_strength = float(acf[local_max_idx])
+
+        if harmonic_strength > 0.3:
+            harmonics.append(
+                {
+                    "harmonic": h,
+                    "lag_samples": local_max_idx,
+                    "strength": harmonic_strength,
+                }
+            )
+
+    return harmonics
+
+
 def find_periodicity(
     trace: WaveformTrace | NDArray[np.floating[Any]],
     *,
@@ -301,32 +487,19 @@ def find_periodicity(
         >>> print(f"Period: {result['period_time']*1e6:.2f} us")
         >>> print(f"Frequency: {result['frequency']/1e3:.1f} kHz")
     """
-    if isinstance(trace, WaveformTrace):
-        data = trace.data
-        fs = trace.metadata.sample_rate
-    else:
-        data = trace
-        if sample_rate is None:
-            raise ValueError("sample_rate required when trace is array")
-        fs = sample_rate
-
+    # Setup: extract data and compute autocorrelation
+    data, fs = _extract_periodicity_data(trace, sample_rate)
     n = len(data)
+    max_period_samples = max_period_samples if max_period_samples is not None else n // 2
 
-    if max_period_samples is None:
-        max_period_samples = n // 2
-
-    # Compute autocorrelation
     _lag_times, acf = autocorrelation(
-        trace,
-        max_lag=max_period_samples,
-        sample_rate=sample_rate if sample_rate else fs,
+        trace, max_lag=max_period_samples, sample_rate=sample_rate if sample_rate else fs
     )
 
-    # Find peaks in autocorrelation (after lag 0)
-    # Look for local maxima
-    acf_search = acf[min_period_samples:]
+    # Processing: find primary peak and harmonics
+    period_samples, strength = _find_primary_peak(acf, min_period_samples)
 
-    if len(acf_search) < 3:
+    if period_samples < 0:
         return {
             "period_samples": np.nan,
             "period_time": np.nan,
@@ -335,46 +508,11 @@ def find_periodicity(
             "harmonics": [],
         }
 
-    # Find local maxima
-    local_max = (acf_search[1:-1] > acf_search[:-2]) & (acf_search[1:-1] > acf_search[2:])
-    max_indices = np.where(local_max)[0] + 1  # +1 for offset from [1:-1]
-
-    if len(max_indices) == 0:
-        # No local maxima found, use global max
-        primary_idx = int(np.argmax(acf_search)) + min_period_samples
-        strength = float(acf[primary_idx])
-    else:
-        # Find strongest peak
-        peak_values = acf_search[max_indices]
-        best_peak_idx = int(np.argmax(peak_values))
-        primary_idx = int(max_indices[best_peak_idx]) + min_period_samples
-        strength = float(acf[primary_idx])
-
-    period_samples = int(primary_idx)
     period_time = period_samples / fs
     frequency = 1.0 / period_time if period_time > 0 else np.nan
+    harmonics = _find_harmonics(acf, period_samples)
 
-    # Find harmonics (peaks at multiples of period)
-    harmonics: list[dict[str, int | float]] = []
-    for h in range(2, 6):  # Check up to 5th harmonic
-        harmonic_lag = h * period_samples
-        if harmonic_lag < len(acf):
-            # Look for peak near expected harmonic
-            search_range = max(1, period_samples // 4)
-            start = int(max(0, harmonic_lag - search_range))
-            end = int(min(len(acf), harmonic_lag + search_range))
-            local_max_idx = int(start + int(np.argmax(acf[start:end])))
-            harmonic_strength = float(acf[local_max_idx])
-
-            if harmonic_strength > 0.3:  # Threshold for significant harmonic
-                harmonics.append(
-                    {
-                        "harmonic": h,
-                        "lag_samples": local_max_idx,
-                        "strength": harmonic_strength,
-                    }
-                )
-
+    # Result building: construct result dictionary
     return {
         "period_samples": period_samples,
         "period_time": float(period_time),
@@ -483,140 +621,198 @@ def correlate_chunked(
         MEM-008: Chunked Correlation
         Oppenheim & Schafer (2009): Discrete-Time Signal Processing, Ch 8
     """
+    _validate_chunked_inputs(signal1, signal2, mode)
+
+    n1, n2 = len(signal1), len(signal2)
+    chunk_size_final = _determine_chunk_size(chunk_size, n1, n2)
+
+    # Use direct correlation for small signals
+    if _should_use_direct_method(n1, n2, chunk_size_final):
+        return _direct_correlate(signal1, signal2, mode)
+
+    # Setup overlap-save parameters
+    params = _setup_overlap_save_params(chunk_size_final, n2)
+    if params is None:
+        return _direct_correlate(signal1, signal2, mode)
+
+    # Prepare kernel FFT and output buffer
+    signal2_flipped = signal2[::-1].copy()
+    kernel_fft = np.fft.fft(signal2_flipped, n=params.nfft)
+    output = np.zeros(_get_output_length(n1, n2, mode), dtype=np.float64)
+
+    # Process signal in chunks
+    _process_chunks_overlap_save(signal1, kernel_fft, output, params, mode)
+
+    return output
+
+
+def _validate_chunked_inputs(
+    signal1: NDArray[np.floating[Any]], signal2: NDArray[np.floating[Any]], mode: str
+) -> None:
+    """Validate inputs for chunked correlation."""
     if len(signal1) == 0 or len(signal2) == 0:
         raise ValueError("Input signals cannot be empty")
-
     if mode not in ("same", "valid", "full"):
         raise ValueError(f"Invalid mode: {mode}. Must be 'same', 'valid', or 'full'")
 
-    n1 = len(signal1)
-    n2 = len(signal2)
-
-    # Determine chunk size
-    if chunk_size is None:
-        # Auto-select: aim for ~100MB chunks
-        bytes_per_sample = 8  # float64
-        target_bytes = 100 * 1024 * 1024
-        chunk_size = min(target_bytes // bytes_per_sample, n1)
-        # Round to power of 2 for FFT efficiency
-        chunk_size = 2 ** int(np.log2(chunk_size))
-
-    # Ensure chunk_size is larger than filter length for overlap-save
-    # Otherwise overlap-save doesn't make sense
+
+def _determine_chunk_size(chunk_size: int | None, n1: int, n2: int) -> int:
+    """Determine optimal chunk size for processing."""
+    if chunk_size is not None:
+        return chunk_size
+
+    # Auto-select: aim for ~100MB chunks (float64 = 8 bytes)
+    target_bytes = 100 * 1024 * 1024
+    auto_chunk = min(target_bytes // 8, n1)
+    log2_val: np.floating[Any] = np.log2(auto_chunk)
+    chunk_power_of_2: int = int(2 ** int(log2_val))  # Round to power of 2
+    return chunk_power_of_2
+
+
+def _should_use_direct_method(n1: int, n2: int, chunk_size: int) -> bool:
+    """Check if direct correlation should be used instead of chunked."""
     min_chunk_size = max(2 * n2, 64)
+    return n1 <= min_chunk_size or n2 >= n1 or chunk_size < min_chunk_size
 
-    # For small signals or when chunk_size is too small, use direct method
-    if n1 <= min_chunk_size or n2 >= n1 or chunk_size < min_chunk_size:
-        mode_literal = cast("Literal['same', 'valid', 'full']", mode)
-        result = np.correlate(signal1, signal2, mode=mode_literal)
-        return result.astype(np.float64)
 
-    # For correlation, we need to flip signal2
-    signal2_flipped = signal2[::-1].copy()
+def _direct_correlate(
+    signal1: NDArray[np.floating[Any]], signal2: NDArray[np.floating[Any]], mode: str
+) -> NDArray[np.float64]:
+    """Perform direct correlation using numpy."""
+    mode_literal = cast("Literal['same', 'valid', 'full']", mode)
+    result = np.correlate(signal1, signal2, mode=mode_literal)
+    return result.astype(np.float64)
+
+
+@dataclass
+class _OverlapSaveParams:
+    """Parameters for overlap-save algorithm."""
 
-    # Overlap-save parameters
-    # L = chunk size, M = filter length
+    chunk_size: int
+    filter_len: int
+    overlap: int
+    step_size: int
+    nfft: int
+
+
+def _setup_overlap_save_params(chunk_size: int, n2: int) -> _OverlapSaveParams | None:
+    """Setup overlap-save algorithm parameters."""
+    min_chunk_size = max(2 * n2, 64)
     L = max(chunk_size, min_chunk_size)
     M = n2
     overlap = M - 1
-
-    # Ensure step size is positive (L must be > overlap)
     step_size = L - overlap
+
     if step_size <= 0:
-        # Fall back to direct method if chunk is too small
-        mode_literal = cast("Literal['same', 'valid', 'full']", mode)
-        result = np.correlate(signal1, signal2, mode=mode_literal)
-        return result.astype(np.float64)
+        return None
 
-    # FFT size (power of 2, >= L + M - 1)
     nfft = int(2 ** np.ceil(np.log2(L + M - 1)))
+    return _OverlapSaveParams(L, M, overlap, step_size, nfft)
 
-    # Pre-compute FFT of flipped signal2 (kernel)
-    kernel_fft = np.fft.fft(signal2_flipped, n=nfft)
 
-    # Output length based on mode
+def _get_output_length(n1: int, n2: int, mode: str) -> int:
+    """Calculate output length based on correlation mode."""
     if mode == "full":
-        output_len = n1 + n2 - 1
+        return n1 + n2 - 1
     elif mode == "same":
-        output_len = n1
+        return n1
     else:  # valid
-        output_len = max(0, n1 - n2 + 1)
+        return max(0, n1 - n2 + 1)
 
-    # Initialize output
-    output = np.zeros(output_len, dtype=np.float64)
 
-    # Process chunks with overlap-save
-    pos = 0  # Position in signal1
-    max_iterations = (n1 // step_size) + 2  # Safety limit
-    iteration = 0
+def _process_chunks_overlap_save(
+    signal1: NDArray[np.floating[Any]],
+    kernel_fft: NDArray[np.complexfloating[Any, Any]],
+    output: NDArray[np.float64],
+    params: _OverlapSaveParams,
+    mode: str,
+) -> None:
+    """Process signal in chunks using overlap-save method."""
+    n1 = len(signal1)
+    pos = 0
+    max_iterations = (n1 // params.step_size) + 2
+
+    for _iteration in range(max_iterations):
+        if pos >= n1:
+            break
+
+        # Extract and process chunk
+        chunk = _extract_chunk(signal1, pos, params, n1)
+        conv_result = _convolve_chunk_fft(chunk, kernel_fft, params.nfft)
+        valid_output = _extract_valid_portion(conv_result, pos, params)
+
+        # Write to output buffer
+        _write_chunk_output(output, valid_output, pos, params, mode)
+
+        pos += params.step_size
+
+
+def _extract_chunk(
+    signal1: NDArray[np.floating[Any]], pos: int, params: _OverlapSaveParams, n1: int
+) -> NDArray[np.floating[Any]]:
+    """Extract chunk with appropriate overlap."""
+    if pos == 0:
+        return signal1[0 : min(params.chunk_size, n1)]
+
+    chunk_start = max(0, pos - params.overlap)
+    chunk_end = min(chunk_start + params.chunk_size, n1)
+    return signal1[chunk_start:chunk_end]
+
+
+def _convolve_chunk_fft(
+    chunk: NDArray[np.floating[Any]],
+    kernel_fft: NDArray[np.complexfloating[Any, Any]],
+    nfft: int,
+) -> NDArray[np.floating[Any]]:
+    """Perform FFT-based convolution on chunk."""
+    chunk_padded = np.zeros(nfft, dtype=np.float64)
+    chunk_padded[: len(chunk)] = chunk
+    chunk_fft = np.fft.fft(chunk_padded)
+    conv_fft = chunk_fft * kernel_fft
+    return np.fft.ifft(conv_fft).real
+
+
+def _extract_valid_portion(
+    conv_result: NDArray[np.floating[Any]], pos: int, params: _OverlapSaveParams
+) -> NDArray[np.floating[Any]]:
+    """Extract valid portion from convolution result."""
+    if pos == 0:
+        valid_start = 0
+        valid_end = min(params.chunk_size, len(conv_result))
+    else:
+        valid_start = params.overlap
+        valid_end = min(params.chunk_size, len(conv_result))
 
-    while pos < n1 and iteration < max_iterations:
-        iteration += 1
+    return conv_result[valid_start:valid_end]
 
-        # Extract chunk with overlap from previous chunk
-        if pos == 0:
-            # First chunk: no overlap needed
-            chunk_start = 0
-            chunk = signal1[0 : min(L, n1)]
-        else:
-            # Subsequent chunks: include overlap
-            chunk_start = max(0, pos - overlap)
-            chunk_end = min(chunk_start + L, n1)
-            chunk = signal1[chunk_start:chunk_end]
-
-        # Zero-pad chunk to FFT size
-        chunk_padded = np.zeros(nfft, dtype=np.float64)
-        chunk_padded[: len(chunk)] = chunk
-
-        # Perform FFT-based convolution
-        chunk_fft = np.fft.fft(chunk_padded)
-        conv_fft = chunk_fft * kernel_fft
-        conv_result = np.fft.ifft(conv_fft).real
-
-        # Extract valid portion (discard transient at start)
-        if pos == 0:
-            # First chunk
-            valid_start = 0
-            valid_end = min(L, len(conv_result))
-        else:
-            # Subsequent chunks: discard overlap region
-            valid_start = overlap
-            valid_end = min(len(chunk), len(conv_result))
-
-        valid_output = conv_result[valid_start:valid_end]
-
-        # Determine output range based on mode
-        if mode == "full":
-            # Full convolution includes all overlap
-            out_start = pos
-            out_end = min(out_start + len(valid_output), output_len)
-        elif mode == "same":
-            # Same mode: center-aligned
-            offset = (M - 1) // 2
-            out_start = max(0, pos - offset)
-            out_end = min(out_start + len(valid_output), output_len)
-            # Adjust valid_output if we're at boundaries
-            if pos == 0 and offset > 0:
-                valid_output = valid_output[offset:]
-        else:  # valid
-            # Valid mode: only where signals fully overlap
-            offset = M - 1
-            if pos < offset:
-                # Skip this chunk, not in valid region yet
-                pos += step_size
-                continue
-            out_start = pos - offset
-            out_end = min(out_start + len(valid_output), output_len)
-
-        # Copy to output
-        copy_len = min(len(valid_output), out_end - out_start)
-        if copy_len > 0:
-            output[out_start : out_start + copy_len] = valid_output[:copy_len]
-
-        # Move to next chunk with guaranteed progress
-        pos += step_size
 
-    return output
+def _write_chunk_output(
+    output: NDArray[np.float64],
+    valid_output: NDArray[np.floating[Any]],
+    pos: int,
+    params: _OverlapSaveParams,
+    mode: str,
+) -> None:
+    """Write chunk output to final buffer based on mode."""
+    output_len = len(output)
+    offset = (params.filter_len - 1) // 2 if mode == "same" else params.filter_len - 1
+
+    if mode == "full":
+        out_start = pos
+    elif mode == "same":
+        out_start = max(0, pos - offset)
+        if pos == 0 and offset > 0:
+            valid_output = valid_output[offset:]
+    else:  # valid
+        if pos < offset:
+            return
+        out_start = pos - offset
+
+    out_end = min(out_start + len(valid_output), output_len)
+    copy_len = min(len(valid_output), out_end - out_start)
+
+    if copy_len > 0:
+        output[out_start : out_start + copy_len] = valid_output[:copy_len]
 
 
 __all__ = [