Faster way to check if elements in numpy array windows are finite

user_2955541 0
Viewed 216

I have a very long NumPy array with 1_000_000_000 elements and I want to slide a 50 element window across the array and ask if all of the elements within the window are finite. If all elements within a 50 element window are all finite then return True (for that window), otherwise, if one or more elements within the 50 element window are not finite then return False (for that window). Continue this assessment until all windows are assessed. A nice way to do this is:

import numpy as np

def rolling_window(a, window):
    a = np.asarray(a)
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)

    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

if __name__ == "__main__":
    a = np.random.rand(100_000_000)  # This is 10x shorter than my real data
    w = 50
    idx = np.random.randint(0, len(a), size=len(a)//10)  # Simulate having np.nan in my array
    a[idx] = np.nan
    print(np.all(rolling_window(np.isfinite(a), w), axis=1))

However, this is slow when my array is of length 1_000_000_000. Is there a faster way to accomplish this that also doesn't require a ton of memory?

numpy arrays python performance
1 Answers

Approach #1: Abuse strided windows directly into the isfinite-mask for the assignment -

def strided_allfinite(a, w):
    m = np.isfinite(a)
    p = rolling_window(m, w)
    nmW = ~m[:w]
    if nmW.any():
        m[:np.flatnonzero(nmW).max()] = False
    p[~m[w-1:]] = False
    return m[:-w+1]

Timings on given sample data :

In [323]: N = 100_000_000
     ...: w = 50
     ...: 
     ...: np.random.seed(0)
     ...: a = np.random.rand(N)  # This is 10x shorter than my real data
     ...: idx = np.random.randint(0, len(a), size=len(a)//10)  # Simulate...
     ...: a[idx] = np.nan

# Original soln
In [324]: %timeit np.all(rolling_window(np.isfinite(a), w), axis=1)
1.61 s ± 14.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [325]: %timeit strided_allfinite(a, w)
556 ms ± 87.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Approach #2

We can leverage convolution -

np.convolve(np.isfinite(a), np.ones(w),'valid')==w

Approach #3

With binary-erosion -

from scipy.ndimage.morphology import binary_erosion

m = np.isfinite(a)
out = binary_erosion(m, np.ones(w, dtype=bool))[w//2:len(a)-w+1+w//2]
user_32938810
Related