TDAmeritrade / stumpy

STUMPY is a powerful and scalable Python library for modern time series analysis

Home Page:https://stumpy.readthedocs.io/en/latest/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Add Faster rolling_isconstant

seanlaw opened this issue · comments

@NimaSarajpoor I came across a numpy "peak-to-peak" function called np.ptp that computes the difference between the max and min. On my machine it is around 2x faster for computing the rolling_isconstant:

@njit(parallel=True, fastmath={"nsz", "arcp", "contract", "afn", "reassoc"})
def _rolling_isconstant(a, w):
    l = a.shape[0] - w + 1
    out = np.empty(l)
    for i in prange(l):
        out[i] = np.ptp(a[i : i + w])

    return np.where(out == 0.0, True, False)


def rolling_isconstant(a, w):
    axis = a.ndim - 1
    return np.apply_along_axis(
            lambda a_row, w: _rolling_isconstant(a_row, w), axis=axis, arr=a, w=w
        )

@seanlaw
Awesome! it is very nice that numba supports this! Thank you for this suggestion :)

In my PC, it is about 6x faster :)

Cool. I will upgrade it to this (and generalize it)