Repository: librosa/librosa

Issue:
Reduce dependency on padding
Many of our functions use padding to simplify vectorization of certain operations, eg:

https://github.com/librosa/librosa/blob/aec87d801144168d95142386a3efd200cb3dde76/librosa/util/utils.py#L943-L946

in `util.localmax` exists solely to make edge cases easy to handle in vectorized code.

This is good because the code is concise, readable, and fast.  But it's also bad because it invokes a full `O(n)` memory copy for what is literally an edge case.

We should conduct a thorough audit of the library for this kind of thing -- I don't think there is very much of it -- and replace pads with more careful code that can operate in place if possiible.  Odds are good that this will require some numbafication if we have to replace vector ops with loops.


Discussion:
#871 has some progress on this
Here's a brief summary of the various occurrences of padding in the code.  These are prioritized by severity.

- stft, probably avoidable #871 : https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/spectrum.py#L224
- Onset envelope calculation.  Possibly avoidable. https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/onset.py#L606-L616
- recurrence→lag in the dense case.  Maybe avoidable by over-allocating up front and then doing an in-place shear.  https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/segment.py#L636
- zero-crossings: maybe avoidable. https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/audio.py#L1174-L1183
- parabolic interpolation, could be optimized out I think: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/pitch.py#L336-L340
- yin, probably not avoidable without substantial hackery: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/pitch.py#L567-L570
- pyin, same: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/pitch.py#L773-L776
- phase vocoding, might be doable: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/spectrum.py#L1257-L1260
- iirt, could be tricky: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/core/spectrum.py#L1401-L1405
- tempogram, could be removed by special-casing the edges: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/feature/rhythm.py#L154-L159
- rms, could be special-cased at the edges: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/feature/spectral.py#L946-L952
- zcr, could be special-cased: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/feature/spectral.py#L1186-L1189
- stack_memory, difficult in general: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/feature/utils.py#L240-L248
- localmax, could be avoided (see OP): https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/util/utils.py#L1023
- localmin, same deal: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/util/utils.py#L1085
- pad_center, unavoidable: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/util/utils.py#L433
- fix_length, unavoidable: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/util/utils.py#L561
- cyclic_gradient, unavoidable: https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/util/utils.py#L1867
- Window construction: unavoidable and not worth modifying https://github.com/librosa/librosa/blob/a3095313a2ba9e37e24f6180c4decc77cd6a2b0d/librosa/filters.py#L426

After a bit of reflection, I think we should remove this from the 0.9.2 milestone.

Getting STFT and friends #871 covered will pick up the vast majority of performance gains here, and the rest would make the code more complex than I'm really comfortable with.  (Though TBH, not very many of these are easy modifications anyway.)

Thinking longer term, I think what would be most helpful here is an abstracted "padded array" class that divides up an ndarray into its padded regions and original region, similar to how a dask array divides up into chunks, where the original region is copy-by-reference.  I don't think this functionality belongs in librosa, but if it made its way into numpy (or a lightweight standalone package that we could depend on), that would be the best option all around IMO.
localmax/localmin, zero crossings, and parabolic interpolation bits could probably be implemented efficiently as [numba stencils](https://numba.readthedocs.io/en/stable/user/stencil.html#numba-stencil).

Onset envelope calculation might also be possible.
Quick prototype of a stenciled and guvectorized (to support arbitrary dimensions) localmax:

```python
@numba.stencil
def _localmax_sten(x):
    return (x[0] > x[-1]) & (x[0] >= x[1])


@numba.guvectorize(['void(int16[:], bool_[:])',
                    'void(int32[:], bool_[:])',
                    'void(int64[:], bool_[:])',
                    'void(float32[:], bool_[:])',
                    'void(float64[:], bool_[:])'], '(n)->(n)')
def _localmax(x, y):
    y[:] = _localmax_sten(x)
    
def localmax(x, *, axis=0):
    
    xi = x.swapaxes(-1, axis)
    lmax = np.empty_like(x, dtype=bool)
    lmaxi = lmax.swapaxes(-1, axis)
    _localmax(xi, lmaxi)
    
    lmaxi[..., -1] = xi[..., -1] > xi[..., -2]
    
    return lmax
```

And a benchmark:

```python
>>> x = np.random.randn(1000, 30000)

>>> # Current implementation in 0.9.2
>>> %memit librosa.util.localmax(x, axis=1)
peak memory: 1108.77 MiB, increment: 245.25 MiB

>>> # Stenciled implementation
>>> %memit localmax(x, axis=1)
peak memory: 890.68 MiB, increment: 0.00 MiB

>>> %timeit librosa.util.localmax(x, axis=1)
271 ms ± 53.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

>>> %timeit localmax(x, axis=1)
48.5 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

---
Some comments on the above code: 

- the swapaxes business is to rotate the target axis around to the end so that numba.guvectorize can parallelize over the leading axes
- I've explicitly coded the supported dtypes here, since this is needed for ufunc dispatch.  I suppose we could add boolean support as well, which would basically turn this into a rising edge detector.
- The above pattern carries over naturally to localmin and zerocrossings (I think).
- Implementing variable lag in the onset strength function could be tricky, since this creates a runtime dependency on the neighborhood size.  I'm not sure how to hack around that yet.
Adding this to the 0.10 docket - I expect what we'll end up with is:

- stft/istft via #1514 
- localmax/localmin via the above

The rest of the identified pad usages seem either low-priority (few code paths touch them) or not worth the added code complexity to optimize.

Local min and max are higher priority because they get used as utility functions, eg within (p)yin and beat tracking.  The optimization is easy to do and the payoff is substantial.

Resolving patch:
diff --git a/librosa/util/utils.py b/librosa/util/utils.py
--- a/librosa/util/utils.py
+++ b/librosa/util/utils.py
@@ -962,6 +962,40 @@ def normalize(S, *, norm=np.inf, axis=0, threshold=None, fill=None):
     return Snorm
 
 
+@numba.stencil
+def _localmax_sten(x):  # pragma: no cover
+    '''Numba stencil for local maxima computation'''
+    return (x[0] > x[-1]) & (x[0] >= x[1])
+
+
+@numba.stencil
+def _localmin_sten(x):  # pragma: no cover
+    '''Numba stencil for local minima computation'''
+    return (x[0] < x[-1]) & (x[0] <= x[1])
+
+
+@numba.guvectorize(['void(int16[:], bool_[:])',
+                    'void(int32[:], bool_[:])',
+                    'void(int64[:], bool_[:])',
+                    'void(float32[:], bool_[:])',
+                    'void(float64[:], bool_[:])'], '(n)->(n)',
+                   cache=True, nopython=True)
+def _localmax(x, y):  # pragma: no cover
+    '''Vectorized wrapper for the localmax stencil'''
+    y[:] = _localmax_sten(x)
+
+
+@numba.guvectorize(['void(int16[:], bool_[:])',
+                    'void(int32[:], bool_[:])',
+                    'void(int64[:], bool_[:])',
+                    'void(float32[:], bool_[:])',
+                    'void(float64[:], bool_[:])'], '(n)->(n)',
+                   cache=True, nopython=True)
+def _localmin(x, y):  # pragma: no cover
+    '''Vectorized wrapper for the localmin stencil'''
+    y[:] = _localmin_sten(x)
+
+
 def localmax(x, *, axis=0):
     """Find local maxima in an array
 
@@ -1007,19 +1041,20 @@ def localmax(x, *, axis=0):
     --------
     localmin
     """
+    # Rotate the target axis to the end
+    xi = x.swapaxes(-1, axis)
 
-    paddings = [(0, 0)] * x.ndim
-    paddings[axis] = (1, 1)
+    # Allocate the output array and rotate target axis
+    lmax = np.empty_like(x, dtype=bool)
+    lmaxi = lmax.swapaxes(-1, axis)
 
-    x_pad = np.pad(x, paddings, mode="edge")
+    # Call the vectorized stencil
+    _localmax(xi, lmaxi)
 
-    inds1 = [slice(None)] * x.ndim
-    inds1[axis] = slice(0, -2)
+    # Handle the edge condition not covered by the stencil
+    lmaxi[..., -1] = xi[..., -1] > xi[..., -2]
 
-    inds2 = [slice(None)] * x.ndim
-    inds2[axis] = slice(2, x_pad.shape[axis])
-
-    return (x > x_pad[tuple(inds1)]) & (x >= x_pad[tuple(inds2)])
+    return lmax
 
 
 def localmin(x, *, axis=0):
@@ -1068,19 +1103,20 @@ def localmin(x, *, axis=0):
     --------
     localmax
     """
+    # Rotate the target axis to the end
+    xi = x.swapaxes(-1, axis)
 
-    paddings = [(0, 0)] * x.ndim
-    paddings[axis] = (1, 1)
-
-    x_pad = np.pad(x, paddings, mode="edge")
+    # Allocate the output array and rotate target axis
+    lmin = np.empty_like(x, dtype=bool)
+    lmini = lmin.swapaxes(-1, axis)
 
-    inds1 = [slice(None)] * x.ndim
-    inds1[axis] = slice(0, -2)
+    # Call the vectorized stencil
+    _localmin(xi, lmini)
 
-    inds2 = [slice(None)] * x.ndim
-    inds2[axis] = slice(2, x_pad.shape[axis])
+    # Handle the edge condition not covered by the stencil
+    lmini[..., -1] = xi[..., -1] < xi[..., -2]
 
-    return (x < x_pad[tuple(inds1)]) & (x <= x_pad[tuple(inds2)])
+    return lmin
 
 
 def peak_pick(x, *, pre_max, post_max, pre_avg, post_avg, delta, wait):


Test patch: