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)