Numpy-векторизованная функция для повторения блоков последовательных элементов

Question

Numpy имеет функцию repeat , которая повторяет каждый элемент массива заданный (на) элемент несколько раз.

Я хочу реализовать функцию, которая выполняет аналогичные функции, но повторяет не отдельные элементы, а блоки различного размера из последовательных элементов. По сути, я хочу следующую функцию:

import numpy as np

def repeat_blocks(a, sizes, repeats):
    b = []    
    start = 0
    for i, size in enumerate(sizes):
        end = start + size
        b.extend([a[start:end]] * repeats[i])
        start = end
    return np.concatenate(b)

Например, учитывая

a = np.arange(20)
sizes = np.array([3, 5, 2, 6, 4])
repeats = np.array([2, 3, 2, 1, 3])

1011 * тогда *

repeat_blocks(a, sizes, repeats)

возвращает

array([ 0,  1,  2, 
        0,  1,  2,

        3,  4,  5,  6,  7, 
        3,  4,  5,  6,  7, 
        3,  4,  5,  6,  7, 

        8,  9, 
        8,  9,

        10, 11, 12, 13, 14, 15,

        16, 17, 18, 19,
        16, 17, 18, 19,
        16, 17, 18, 19 ])

Я хочу запихнуть эти петли во имя исполнения. Это возможно? Если да, то как?

Answer 1

Эта функция является отличным кандидатом для ускорения с помощью Numba:

@numba.njit
def repeat_blocks_jit(a, sizes, repeats):
    out = np.empty((sizes * repeats).sum(), a.dtype)
    start = 0
    oi = 0
    for i, size in enumerate(sizes):
        end = start + size
        for rep in range(repeats[i]):
            oe = oi + size
            out[oi:oe] = a[start:end]
            oi = oe
        start = end
    return out

Это значительно быстрее, чем чистое решение Divakar NumPy, и намного ближе к исходному коду. Я не приложил никаких усилий, чтобы оптимизировать его. Обратите внимание, что np.dot() и np.repeat() здесь нельзя использовать, но это не имеет значения, когда весь код скомпилирован.

Кроме того, поскольку njit означает режим "nopython", вы даже можете использовать @numba.njit(nogil=True) и получить многоядерное ускорение, если вам нужно выполнить много таких вызовов.

Answer 2

Вот один векторизованный подход с использованием cumsum -

# Get repeats for each group using group lengths/sizes
r1 = np.repeat(np.arange(len(sizes)), repeats)

# Get total size of output array, as needed to initialize output indexing array
N = (sizes*repeats).sum() # or np.dot(sizes, repeats)

# Initialize indexing array with ones as we need to setup incremental indexing
# within each group when cumulatively summed at the final stage. 
# Two steps here:
# 1. Within each group, we have multiple sequences, so setup the offsetting
# at each sequence lengths by the seq. lengths preceeeding those.
id_ar = np.ones(N, dtype=int)
id_ar[0] = 0
insert_index = sizes[r1[:-1]].cumsum()
insert_val = (1-sizes)[r1[:-1]]

# 2. For each group, make sure the indexing starts from the next group's
# first element. So, simply assign 1s there.
insert_val[r1[1:] != r1[:-1]] = 1

# Assign index-offseting values
id_ar[insert_index] = insert_val

# Finally index into input array for the group repeated o/p
out = a[id_ar.cumsum()]