Как векторизовать это для цикла?

Question

У меня есть массив numy f с длиной n и матрица numpy A с размером n x m . Я хочу разбить f и A на r штук f1 , ..., fr и A1 , ..., Ar , затем выполните вычисления fi * Ai (умножение матрицы на вектор x в математическом смысле) с каждым fi являющийся вектором строк с числом столбцов, равным количеству строк Ai . Результатом будет вектор строки 1 x m . Идея состоит в том, чтобы соединить все эти векторы строк, чтобы сформировать матрицу B = [[f1 * A1], [f2 * A2], ..., [fr * Ar]] (обратите внимание, что это будет матрица размером r x m ).

Предположим, что f и A уже определены. Также предположим, что соответствующий индекс фигур находится в списке [0, d1, ... dr] . Например, f1 = f [d [0]: d [1]] и f2 = f [d [1]: d [2]] ). Я использовал следующий фрагмент кода для решения моей проблемы:

B = numpy.zeros([r,m])
for i in range(0,r):
    lower = d[i]
    upper = d[i+1]
    B[i,:] = f[lower:upper].dot(A[lower:upper,:])

Проблема в том, что этот фрагмент кода будет вычисляться несколько раз в моей программе. Я слышал раньше, что цикл для Python медленный, и на самом деле узким местом моего кода является эта часть. Я не могу понять, как это векторизовать, но я чувствую, что это возможно. Я надеялся, что кто-то здесь сможет показать мне путь. Спасибо.

Answer 1

Вы можете использовать np.add.reduceat:

# example data
>>> f = np.arange(10)
>>> A = np.arange(50).reshape(10, 5)
>>> split = [0, 3, 5, 10]
>>> 
# reduceat
>>> np.add.reduceat(f[:, None] * A, split[:-1], axis=0)
array([[  25,   28,   31,   34,   37],
       [ 125,  132,  139,  146,  153],
       [1275, 1310, 1345, 1380, 1415]])
>>> 
# double check against list comprehension
>>> [fi @ Ai for fi, Ai in zip(*map(np.split, (f, A), 2*(split[1:-1],)))]
[array([25, 28, 31, 34, 37]), array([125, 132, 139, 146, 153]), array([1275, 1310, 1345, 1380, 1415])]

Я не удивлюсь, если бы понимание списка, решение @ hpaulj или цикл OP были быстрее из-за blas ускоренного умножения матриц.

Answer 2

Я предполагаю, что это действительный MCVE:

In [139]: f = np.arange(10)
In [140]: A = np.arange(20).reshape(10,2)
In [141]: f.dot(A)
Out[141]: array([570, 615])
In [142]: d = [0,2,5,10]
In [143]: for i,j in zip(d[:-1],d[1:]):
     ...:     print(f[i:j].dot(A[i:j,:]))
     ...:     
[2 3]
[58 67]
[510 545]

, где 570 = 2+58+510.

In [145]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[145]: 
array([[  2,   3],
       [ 58,  67],
       [510, 545]])

Учитывая, что срезы i:j могут различаться по длине, может быть трудно «векторизовать» это в истинном смысле. Мы можем скрыть итерации, но записать их так, чтобы переместить все итерации в скомпилированный код, будет сложно. Накопительные операции, такие как cumsum, часто являются лучшим выбором. Нам часто приходится отступать назад и смотреть на проблему с другой точки зрения (в отличие от простого удаления цикла).

numba и cython часто используются для ускорения итеративных решений, но я не буду вдаваться в них.

---

Если d делит массивы на равные части, мы можем использовать изменение формы для вычисления частей:

In [228]: A.shape
Out[228]: (10, 2)
In [229]: f.shape
Out[229]: (10,)
In [230]: f2 = f.reshape(2,5)
In [231]: A2 = A.reshape(2,5,2)

In [233]: np.einsum('ij,ijk->ik',f2,A2)
Out[233]: 
array([[ 60,  70],
       [510, 545]])

Оператор matmul также работает, хотя требует некоторого изменения размеров:

In [236]: (f2[:,None,:]@A2)[:,0,:]
Out[236]: 
array([[ 60,  70],
       [510, 545]])

---

Если d делит массивы только на пару размеров, я думаю, мы могли бы сгруппировать общие размеры и выполнить приведенные выше изменения и einsum для каждой группы, но я не проработал детали:

In [238]: d = [0,2,5,7,10]
In [239]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[239]: 
array([[  2,   3],
       [ 58,  67],
       [122, 133],
       [388, 412]])
In [240]: [f[i:j] for i,j in zip(d[:-1],d[1:])]
Out[240]: [array([0, 1]), array([2, 3, 4]), array([5, 6]), array([7, 8, 9])]

Здесь у нас есть 2 группы, одна из которых имеет длину 2, а другая - длину 3.