Is there a faster alternative to np.diff?

Viewed 1502

I'm concerned with the speed of the following function:

def cch(tau):
    return np.sum(abs(-1*np.diff(cartprod)-tau)<0.001)

Where "cartprod" is a variable for a list that looks like this:

cartprod = np.ndarray([[0.0123,0.0123],[0.0123,0.0459],...])

The length of this list is about 25 million. Basically, I'm trying to find a significantly faster way to return a list of differences for every pair list in that np.ndarray. Is there an algorithmic way or function that's faster than np.diff? Or, is np.diff the end all be all? I'm also open to anything else.

EDIT: Thank you all for your solutions!

3 Answers

We can leverage multi-core with numexpr module for large data and to gain memory efficiency and hence performance with some help from array-slicing -

import numexpr as ne

def cch_numexpr(a, tau):
    d = {'a0':a[:,0],'a1':a[:,1]}
    return np.count_nonzero(ne.evaluate('abs(a0-a1-tau)<0.001',d))

Sample run and timings on 25M sized data -

In [83]: cartprod = np.random.rand(25000000,2)

In [84]: cch(cartprod, tau=0.5) == cch_numexpr(cartprod, tau=0.5)
Out[84]: True

In [85]: %timeit cch(cartprod, tau=0.5)
10 loops, best of 3: 150 ms per loop

In [86]: %timeit cch_numexpr(cartprod, tau=0.5)
10 loops, best of 3: 25.5 ms per loop

Around 6x speedup.

This was with 8 threads. Thus, with more number of threads available for compute, it should improve further. Related post on how to control multi-core functionality.

I think you're hitting a wall by repeatedly returning multiple np.arrays of length ~25 million rather than np.diff being slow. I wrote an equivalent function that iterates over the array and tallies the results as it goes along. The function needs to be jitted with numba to be fast. I hope that is acceptable.

arr = np.random.rand(25000000, 2)

def cch(tau, cartprod):
    return np.sum(abs(-1*np.diff(cartprod)-tau)<0.001)
%timeit cch(0.01, arr)

@jit(nopython=True)
def cch_jit(tau, cartprod):
    count = 0
    tau = -tau
    for i in range(cartprod.shape[0]):
        count += np.less(np.abs(tau - (cartprod[i, 1]- cartprod[i, 0])), 0.001)
    return count
%timeit cch_jit(0.01, arr)

produces

294 ms ± 2.82 ms 
42.7 ms ± 483 µs 

which is about ~6 times faster.

Just out of curiosity I compared the solutions of @Divakar numexpr and @alexdor numba.jit. The implementation numexpr.evaluate seems to be twice as fast as using numba's jit compiler. The results are shown for 100 runs each:

np.sum:          111.07543396949768
numexpr:         12.282189846038818
JIT:             6.2505223751068115
'np.sum' returns same result as 'numexpr'
'np.sum' returns same result as 'jit'
'numexpr' returns same result as 'jit'

Script so reproduce the results:

import numpy as np
import time
import numba
import numexpr

arr = np.random.rand(25000000, 2)
runs = 100

def cch(tau, cartprod):
    return np.sum(abs(-1*np.diff(cartprod)-tau)<0.001)

def cch_ne(tau, cartprod):
    d = {'a0':cartprod[:,0],'a1':cartprod[:,1], 'tau': tau}
    count = np.count_nonzero(numexpr.evaluate('abs(a0-a1-tau)<0.001',d))
    return count

@numba.jit(nopython=True)
def cch_jit(tau, cartprod):
    count = 0
    tau = -tau
    for i in range(cartprod.shape[0]):
        count += np.less(np.abs(tau - (cartprod[i, 1]- cartprod[i, 0])), 0.001)
    return count

start = time.time()
for x in range(runs):
    x1 = cch(0.01, arr)
print('np.sum:\t\t', time.time() - start)

start = time.time()
for x in range(runs):
    x2 = cch_ne(0.01, arr)
print('numexpr:\t', time.time() - start)

x3 = cch_jit(0.01, arr)
start = time.time()
for x in range(runs):
    x3 = cch_jit(0.01, arr)
print('JIT:\t\t', time.time() - start)

if x1 == x2: print('\'np.sum\' returns same result as \'numexpr\'')
if x1 == x3: print('\'np.sum\' returns same result as \'jit\'')
if x2 == x3: print('\'numexpr\' returns same result as \'jit\'')
Related