CODEWITHSUNDEEP
pythonarraysperformancenumpyvectorization
Arty
Arty

Reputation: 16737

Is there multi arange in NumPy?

Numpy's arange accepts only single scalar values for start/stop/step. Is there a multi version of this function? Which can accept array inputs for start/stop/step? E.g. having input 2D array like:

[[1 5 1], # start/stop/step first
 [3 8 2]] # start/stop/step second

should create array consisting of concatenation of aranges for every row of input (each start/stop/step), input above should create 1D array

1 2 3 4 3 5 7

i.e. we need to design such function that it does next:

print(np.multi_arange(np.array([[1,5,1],[3,8,2]])))
# prints:
# array([1, 2, 3, 4, 3, 5, 7])

And this function should be efficient (pure numpy), i.e. very fast process input array of shape (10000, 3) without pure-Python looping.

Of cause it is possible to create pure Python's loop (or listcomp) to create arange for each row and concatenate results of this loop. But I have very many rows with triples start/stop/step and need to have efficient and fast code, hence looking for pure numpy function.

Why do I need it. I needed this for several tasks. One of this is for indexing - suppose I have 1D array a and I need to extract many (possibly intersecting) subranges of this array. If I had that multi version of arange I would just do:

values = a[np.multi_arange(starts_stops_steps)]

Maybe it is possible to create multi arange function using some combinations of numpy functions? Can you suggest?

Also maybe there are some more efficient solutions for the specific case of extracting subranges of 1D array (see last line of code above) without creating all indexes using multi_arange?

Upvotes: 1

Views: 911

Answers (3)

Arty
Arty

Reputation: 16737

I've just came up with my solution using numba. Still I prefer numpy-only solution if we find best one not to carry heavy numba JIT compiler.

I've also tested @Divakar solution in my code.

Next code output is:

naive_multi_arange 0.76601 sec
arty_multi_arange 0.01801 sec 42.52 speedup
divakar_multi_arange 0.05504 sec 13.92 speedup

Meaning my numba solution has 42x speedup, @Divakar's numpy solution has 14x speedup.

Next code can be also run online here.

import time, random
import numpy as np, numba

@numba.jit(nopython = True)
def arty_multi_arange(a):
    starts, stops, steps = a[:, 0], a[:, 1], a[:, 2]
    pos = 0
    cnt = np.sum((stops - starts + steps - np.sign(steps)) // steps, dtype = np.int64)
    res = np.zeros((cnt,), dtype = np.int64)
    for i in range(starts.size):
        v, stop, step = starts[i], stops[i], steps[i]
        if step > 0:
            while v < stop:
                res[pos] = v
                pos += 1
                v += step
        elif step < 0:
            while v > stop:
                res[pos] = v
                pos += 1
                v += step
    assert pos == cnt
    return res
    
def divakar_multi_arange(a):
    steps = a[:,2]
    lens = ((a[:,1]-a[:,0]) + steps-np.sign(steps))//steps
    b = np.repeat(steps, lens)
    ends = (lens-1)*steps + a[:,0]
    b[0] = a[0,0]
    b[lens[:-1].cumsum()] = a[1:,0] - ends[:-1]
    return b.cumsum()
    
random.seed(0)
neg_prob = 0.5
N = 100000
minv, maxv, maxstep = -100, 300, 15
steps = [random.randrange(1, maxstep + 1) * ((1, -1)[random.random() < neg_prob]) for i in range(N)]
starts = [random.randrange(minv + 1, maxv) for i in range(N)]
stops = [random.randrange(*(((starts[i] + 1, maxv + 1), (minv, starts[i]))[steps[i] < 0])) for i in range(N)]
joined = np.array([starts, stops, steps], dtype = np.int64).T

tb = time.time()
aref = np.concatenate([np.arange(joined[i, 0], joined[i, 1], joined[i, 2], dtype = np.int64) for i in range(N)])
npt = time.time() - tb
print('naive_multi_arange', round(npt, 5), 'sec')

for func in ['arty_multi_arange', 'divakar_multi_arange']:
    globals()[func](joined)
    tb = time.time()
    a = globals()[func](joined)
    myt = time.time() - tb
    print(func, round(myt, 5), 'sec', round(npt / myt, 2), 'speedup')
    assert a.size == aref.size, (a.size, aref.size)
    assert np.all(a == aref), np.vstack((np.flatnonzero(a != aref)[:5], a[a != aref][:5], aref[a != aref][:5])).T

Upvotes: 0

Divakar
Divakar

Reputation: 221504

Here's a vectorized one with cumsum that accounts for positive and negative stepsizes -

def multi_arange(a):
    steps = a[:,2]
    lens = ((a[:,1]-a[:,0]) + steps-np.sign(steps))//steps
    b = np.repeat(steps, lens)
    ends = (lens-1)*steps + a[:,0]
    b[0] = a[0,0]
    b[lens[:-1].cumsum()] = a[1:,0] - ends[:-1]
    return b.cumsum()

If you need to validate for valid ranges : (start < stop when step > 0) and (start > stop when step < 0) , use a pre-processing step :

a = a[((a[:,1] > a[:,0]) & (a[:,2]>0) | (a[:,1] < a[:,0]) & (a[:,2]<0))]

Sample run -

In [17]: a
Out[17]: 
array([[ 1,  5,  1],
       [ 3,  8,  2],
       [18,  6, -2]])

In [18]: multi_arange(a)
Out[18]: array([ 1,  2,  3,  4,  3,  5,  7, 18, 16, 14, 12, 10,  8])

Upvotes: 1

hpaulj
hpaulj

Reputation: 231325

In [1]: np.r_[1:5:1, 3:8:2]
Out[1]: array([1, 2, 3, 4, 3, 5, 7])

In [2]: np.hstack((np.arange(1,5,1),np.arange(3,8,2)))
Out[2]: array([1, 2, 3, 4, 3, 5, 7])

The r_ version is nice and compact, but not faster:

In [3]: timeit np.r_[1:5:1, 3:8:2]
23.9 µs ± 34.6 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [4]: timeit np.hstack((np.arange(1,5,1),np.arange(3,8,2)))
11.2 µs ± 19.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Upvotes: 0

Related Questions