Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save makoc/42b8653070058208618b to your computer and use it in GitHub Desktop.
Save makoc/42b8653070058208618b to your computer and use it in GitHub Desktop.
Numpy tricks I stumbled upon
"""Impressive little thing how einsum + stride_tricks can beat numpys build in C functions
for correlate (for large data). (ok depending on the implementation of np.correlate, the
comparison is not fair, but still rather impressive that you can get comparable speeds)
"""
import numpy as np
import stride_tricks as st # stride_tricks.py gist
a = np.random.random((100,100,100)).ravel()
stamp = np.random.random((3,3,3)).ravel()
np.correlate(a[:50000], stamp.ravel(), mode='valid')
# is quivalent to:
np.einsum('ab,b->a', st.rolling_window(a[:50000], 27), stamp.ravel())
# and actually _outperforms_ np.correlate for this size (and larger sizes)
# Plus it can be generalized to higher dimensions:
a = np.random.random((100,100,100))
stamp = np.random.random((3,3,3))
# the ReIndex and .copy() here is so that the way the iteration is done is optimized for
# the CPU. I actually am not sure why this is the fastest option on my cpu:
np.einsum('ijkabc,abc->ijk', st.ReIndex(st.rolling_window(a, stamp.shape))[...,-1,-2,-3],
st.ReIndex(stamp)[-1,-2,-3].copy())
# which quite simply vastly outperforms scipy.signal.correlate and is still only
# approximatly 3 times slower then scipy.ndimage.correlate (which is not equivalent though)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment