Как эффективно объединить многие вызовы arange в numpy?

Я хотел бы векторизовать вызовы, такие как numpy.arange(0, cnt_i) за вектором cnt значения и объединить результаты, как этот фрагмент:

import numpy
cnts = [1,2,3]
numpy.concatenate([numpy.arange(cnt) for cnt in cnts])

array([0, 0, 1, 0, 1, 2])

к сожалению, приведенный выше код очень неэффективен из-за временных массивов и цикла понимания списка.

есть ли способ сделать это более эффективно в numpy?

6
arrays numpy python vectorization

3 ответов

вот полностью векторизовать функции:

def multirange(counts):
    counts = np.asarray(counts)
    # Remove the following line if counts is always strictly positive.
    counts = counts[counts != 0]

    counts1 = counts[:-1]
    reset_index = np.cumsum(counts1)

    incr = np.ones(counts.sum(), dtype=int)
    incr[0] = 0
    incr[reset_index] = 1 - counts1

    # Reuse the incr array for the final result.
    incr.cumsum(out=incr)
    return incr

вот вариант ответа @Developer, который вызывает только arange после:

def multirange_loop(counts):
    counts = np.asarray(counts)
    ranges = np.empty(counts.sum(), dtype=int)
    seq = np.arange(counts.max())
    starts = np.zeros(len(counts), dtype=int)
    starts[1:] = np.cumsum(counts[:-1])
    for start, count in zip(starts, counts):
        ranges[start:start + count] = seq[:count]
    return ranges

и вот оригинальная версия, написанная как функция:

def multirange_original(counts):
    ranges = np.concatenate([np.arange(count) for count in counts])
    return ranges

демо:

In [296]: multirange_original([1,2,3])
Out[296]: array([0, 0, 1, 0, 1, 2])

In [297]: multirange_loop([1,2,3])
Out[297]: array([0, 0, 1, 0, 1, 2])

In [298]: multirange([1,2,3])
Out[298]: array([0, 0, 1, 0, 1, 2])

сравните время, используя больший массив подсчетов:

In [299]: counts = np.random.randint(1, 50, size=50)

In [300]: %timeit multirange_original(counts)
10000 loops, best of 3: 114 µs per loop

In [301]: %timeit multirange_loop(counts)
10000 loops, best of 3: 76.2 µs per loop

In [302]: %timeit multirange(counts)
10000 loops, best of 3: 26.4 µs per loop
4

попробуйте сделать следующее для решения проблемы с памятью, эффективность почти такая же.

out = np.empty((sum(cnts)))
k = 0
for cnt in cnts:
    out[k:k+cnt] = np.arange(cnt)
    k += cnt

поэтому конкатенация не используется.

3

np.tril_indices в значительной степени делает это для вас:

In [28]: def f(c):
   ....:     return np.tril_indices(c, -1)[1]

In [29]: f(10)
Out[29]:
array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 0, 1,
       2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8])

In [33]: %timeit multirange(range(10))
10000 loops, best of 3: 93.2 us per loop

In [34]: %timeit f(10)
10000 loops, best of 3: 68.5 us per loop

гораздо быстрее, чем @Warren Weckesser multirange когда размер мал.

но становится намного медленнее, когда размер больше (@hpaulj, у вас есть очень хорошая точка):

In [36]: %timeit multirange(range(1000))
100 loops, best of 3: 5.62 ms per loop

In [37]: %timeit f(1000)
10 loops, best of 3: 68.6 ms per loop
1