Repository: librosa/librosa

Issue:
Accelerate parabolic interpolation in yin/pyin
**Is your feature request related to a problem? Please describe.**

In the general theme of the 0.10 milestone, I noticed a step in yin/pyin that could be done much more efficiently.  The parabolic interpolation function:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L427-L450
works by slicing the input array and doing some vectorized arithmetic ops to compute the intermediate values.  Each intermediate value (`parabola_a` and `parabola_b`) will therefore require allocating almost as much space as the input array (`y_frames`) - so we're tripling memory usage unnecessarily.


**Describe the solution you'd like**

The operation here is entirely local, and could be easily implemented by a numba stencil operation. This is very similar to what we did for localmax/localmin in #1533,  eg:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/util/utils.py#L983-L1003


**Describe alternatives you've considered**
This isn't the biggest speed bottleneck in yin/pyin, but it would be easy to improve so we should do it anyway.

Accelerate parabolic interpolation in yin/pyin
**Is your feature request related to a problem? Please describe.**

In the general theme of the 0.10 milestone, I noticed a step in yin/pyin that could be done much more efficiently.  The parabolic interpolation function:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L427-L450
works by slicing the input array and doing some vectorized arithmetic ops to compute the intermediate values.  Each intermediate value (`parabola_a` and `parabola_b`) will therefore require allocating almost as much space as the input array (`y_frames`) - so we're tripling memory usage unnecessarily.


**Describe the solution you'd like**

The operation here is entirely local, and could be easily implemented by a numba stencil operation. This is very similar to what we did for localmax/localmin in #1533,  eg:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/util/utils.py#L983-L1003


**Describe alternatives you've considered**
This isn't the biggest speed bottleneck in yin/pyin, but it would be easy to improve so we should do it anyway.



Discussion:
A few minutes of hacking later...

```python
@numba.stencil
def _parabol_sten(x):
    a = x[1] + x[-1] - 2 * x[0]
    b = (x[1] - x[-1]) / 2
    
    if np.abs(b) >= np.abs(a):
        return 0
    
    return -b / a


@numba.guvectorize(['void(float32[:], float32[:])',
                    'void(float64[:], float64[:])'], '(n)->(n)',
                   cache=True, nopython=True)
def _parabol(x, y):  # pragma: no cover
    '''Vectorized wrapper for the parabolic interpolation stencil'''
    y[:] = _parabol_sten(x)
    
def parabolic_interpolation(x, *, axis=-2):
    # Rotate the target axis to the end
    xi = x.swapaxes(-1, axis)

    # Allocate the output array and rotate target axis
    lp = np.empty_like(x)
    lpi = lp.swapaxes(-1, axis)

    # Call the vectorized stencil
    _parabol(xi, lpi)

    # Handle the edge condition not covered by the stencil
    lpi[..., -1] = 0
    lpi[..., 0] = 0

    return lp
```

Comparison:

```python
>>> y, sr = librosa.load(librosa.ex('brahms', hq=True), sr=None)
>>> yf = librosa.util.frame(y, frame_length=2048, hop_length=512)
>>> lp = librosa.core.pitch._parabolic_interpolation(yf)
>>> lp2 = parabolic_interpolation(yf)
>>> np.allclose(lp2, lp)
True
>>> %timeit lp = librosa.core.pitch._parabolic_interpolation(yf)
319 ms ± 17.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit lp2 = parabolic_interpolation(yf)
39.9 ms ± 1.89 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

~~Needs a bit of tidying up, eg to handle numerical stability in the division, but the speedup seems good to me.~~ EDIT: updated to avoid division stability issues and simplify operations.  I think this is solid now.
We could probably also refactor `piptrack` to make use of this acceleration:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L321-L330
though it would need to be expanded or supplemented to expose the magnitude skew:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L338
Probably this should just be a separate helper function.
A few minutes of hacking later...

```python
@numba.stencil
def _parabol_sten(x):
    a = x[1] + x[-1] - 2 * x[0]
    b = (x[1] - x[-1]) / 2
    
    if np.abs(b) >= np.abs(a):
        return 0
    
    return -b / a


@numba.guvectorize(['void(float32[:], float32[:])',
                    'void(float64[:], float64[:])'], '(n)->(n)',
                   cache=True, nopython=True)
def _parabol(x, y):  # pragma: no cover
    '''Vectorized wrapper for the parabolic interpolation stencil'''
    y[:] = _parabol_sten(x)
    
def parabolic_interpolation(x, *, axis=-2):
    # Rotate the target axis to the end
    xi = x.swapaxes(-1, axis)

    # Allocate the output array and rotate target axis
    lp = np.empty_like(x)
    lpi = lp.swapaxes(-1, axis)

    # Call the vectorized stencil
    _parabol(xi, lpi)

    # Handle the edge condition not covered by the stencil
    lpi[..., -1] = 0
    lpi[..., 0] = 0

    return lp
```

Comparison:

```python
>>> y, sr = librosa.load(librosa.ex('brahms', hq=True), sr=None)
>>> yf = librosa.util.frame(y, frame_length=2048, hop_length=512)
>>> lp = librosa.core.pitch._parabolic_interpolation(yf)
>>> lp2 = parabolic_interpolation(yf)
>>> np.allclose(lp2, lp)
True
>>> %timeit lp = librosa.core.pitch._parabolic_interpolation(yf)
319 ms ± 17.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit lp2 = parabolic_interpolation(yf)
39.9 ms ± 1.89 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

~~Needs a bit of tidying up, eg to handle numerical stability in the division, but the speedup seems good to me.~~ EDIT: updated to avoid division stability issues and simplify operations.  I think this is solid now.
We could probably also refactor `piptrack` to make use of this acceleration:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L321-L330
though it would need to be expanded or supplemented to expose the magnitude skew:
https://github.com/librosa/librosa/blob/0a60e0a87888d1bc8b2c6098e17d2b98536a28a0/librosa/core/pitch.py#L338
Probably this should just be a separate helper function.

Resolving patch:
diff --git a/librosa/core/pitch.py b/librosa/core/pitch.py
--- a/librosa/core/pitch.py
+++ b/librosa/core/pitch.py
@@ -5,6 +5,7 @@
 import warnings
 import numpy as np
 import scipy
+import numba
 
 
 from .spectrum import _spectrogram
@@ -326,20 +327,9 @@ def piptrack(
     # Do the parabolic interpolation everywhere,
     # then figure out where the peaks are
     # then restrict to the feasible range (fmin:fmax)
-    avg = 0.5 * (S[..., 2:, :] - S[..., :-2, :])
-
-    shift = 2 * S[..., 1:-1, :] - S[..., 2:, :] - S[..., :-2, :]
-
-    # Suppress divide-by-zeros.
-    # Points where shift == 0 will never be selected by localmax anyway
-    shift = avg / (shift + (np.abs(shift) < util.tiny(shift)))
-
-    # Pad back up to the same shape as S
-    padding = [(0, 0) for _ in S.shape]
-    padding[-2] = (1, 1)
-    avg = np.pad(avg, padding, mode="constant")
-    shift = np.pad(shift, padding, mode="constant")
-
+    avg = np.gradient(S, axis=-2)
+    shift = _parabolic_interpolation(S, axis=-2)
+    # this will get us the interpolated peak value
     dskew = 0.5 * avg * shift
 
     # Pre-allocate output
@@ -358,6 +348,7 @@ def piptrack(
     if callable(ref):
         ref_value = threshold * ref(S, axis=-2)
         # Reinsert the frequency axis here, in case the callable doesn't
+
         # support keepdims=True
         ref_value = np.expand_dims(ref_value, -2)
     else:
@@ -433,30 +424,61 @@ def _cumulative_mean_normalized_difference(
     return yin_frames
 
 
-def _parabolic_interpolation(y_frames: np.ndarray) -> np.ndarray:
+@numba.stencil
+def _pi_stencil(x: np.ndarray) -> np.ndarray:
+    '''Stencil to compute local parabolic interpolation'''
+
+    a = x[1] + x[-1] - 2 * x[0]
+    b = (x[1] - x[-1]) / 2
+
+    if np.abs(b) >= np.abs(a):
+        # If this happens, we'll shift by more than 1 bin
+        return 0
+
+    return -b / a
+
+
+@numba.guvectorize(['void(float32[:], float32[:])',
+                    'void(float64[:], float64[:])'], '(n)->(n)',
+                   cache=True, nopython=True)
+def _pi_wrapper(x: np.ndarray, y: np.ndarray) -> None:  # pragma: no cover
+    '''Vectorized wrapper for the parabolic interpolation stencil'''
+    y[:] = _pi_stencil(x)
+
+
+def _parabolic_interpolation(x: np.ndarray, *, axis: int=-2) -> np.ndarray:
     """Piecewise parabolic interpolation for yin and pyin.
 
     Parameters
     ----------
-    y_frames : np.ndarray [shape=(frame_length, n_frames)]
-        framed audio time series.
+    x : np.ndarray
+        array to interpolate
+    axis : int
+        axis along which to interpolate
 
     Returns
     -------
-    parabolic_shifts : np.ndarray [shape=(frame_length, n_frames)]
-        position of the parabola optima
+    parabolic_shifts : np.ndarray [shape=x.shape]
+        position of the parabola optima (relative to bin indices)
+
+        Note: the shift at bin `n` is determined as 0 if the estimated
+        optimum is outside the range `[n-1, n+1]`.
     """
+    # Rotate the target axis to the end
+    xi = x.swapaxes(-1, axis)
 
-    parabolic_shifts = np.zeros_like(y_frames)
-    parabola_a = (
-        y_frames[..., :-2, :] + y_frames[..., 2:, :] - 2 * y_frames[..., 1:-1, :]
-    ) / 2
-    parabola_b = (y_frames[..., 2:, :] - y_frames[..., :-2, :]) / 2
-    parabolic_shifts[..., 1:-1, :] = -parabola_b / (
-        2 * parabola_a + util.tiny(parabola_a)
-    )
-    parabolic_shifts[np.abs(parabolic_shifts) > 1] = 0
-    return parabolic_shifts
+    # Allocate the output array and rotate target axis
+    shifts = np.empty_like(x)
+    shiftsi = shifts.swapaxes(-1, axis)
+
+    # Call the vectorized stencil
+    _pi_wrapper(xi, shiftsi)
+
+    # Handle the edge condition not covered by the stencil
+    shiftsi[..., -1] = 0
+    shiftsi[..., 0] = 0
+
+    return shifts
 
 
 def yin(


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1156,7 +1156,7 @@ def test_piptrack_errors():
     pitches, mags = librosa.piptrack(
         y=None,
         sr=22050,
-        S=np.asarray([[1, 0, 0]]).T,
+        S=np.asarray([[1.0, 0.0, 0.0]]).T,
         n_fft=4096,
         hop_length=None,
         fmin=150,