Как я могу векторизовать этот тройной -цикл над двумерными массивами в numpy?

Могу ли я исключить все циклы Python в этом вычислении:

result[i,j,k] = (x[i] * y[j] * z[k]).sum()

где x[i], y[j], z[k]— векторы длины N, а x, y, zимеют первые измерения длины A, B, Cст. вывод - форма (A,B,C), и каждый элемент сумма тройного -произведения (элемента -мудрого ).

Я могу уменьшить его с 3 до 1 цикла (код ниже ), но застрял, пытаясь убрать последнюю петлю.

При необходимости могу сделатьA=B=C(через небольшое количество заполнения ).

# Example with 3 loops, 2 loops, 1 loop (testing omitted)

N = 100 # more like 100k in real problem
A =   2 # more like 20 in real problem
B =   3 # more like 20 in real problem
C =   4 # more like 20 in real problem

import numpy
x = numpy.random.rand(A, N)
y = numpy.random.rand(B, N)
z = numpy.random.rand(C, N)

# outputs of each variant
result_slow = numpy.empty((A,B,C))
result_vec_C = numpy.empty((A,B,C))
result_vec_CB = numpy.empty((A,B,C))

# 3 nested loops
for i in range(A):
    for j in range(B):
        for k in range(C):
            result_slow[i,j,k] = (x[i] * y[j] * z[k]).sum()

# vectorize loop over C (2 nested loops)
for i in range(A):
    for j in range(B):
        result_vec_C[i,j,:] = (x[i] * y[j] * z).sum(axis=1)

# vectorize one C and B (one loop)
for i in range(A):
    result_vec_CB[i,:,:] = numpy.dot(x[i] * y, z.transpose())

numpy.testing.assert_almost_equal(result_slow, result_vec_C)
numpy.testing.assert_almost_equal(result_slow, result_vec_CB)
7
задан Joseph Hastings 29 June 2012 в 16:23
поделиться