Làm thế nào tôi có thể vector hóa vòng lặp ba vòng trên 2d mảng này trong không?

Question

Tôi có thể loại bỏ tất cả các vòng Python trong tính toán này:

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

nơi x[i], y[j], z[k] là vectơ có độ dài N và x, y, z có kích thước đầu tiên với chiều dài A, B, C s.t. đầu ra là hình dạng (A,B,C) và mỗi phần tử là tổng của một sản phẩm ba (yếu tố khôn ngoan).

Tôi có thể gỡ xuống từ 3 đến 1 vòng (mã bên dưới), nhưng bị kẹt cố gắng để loại bỏ vòng lặp cuối cùng.

Nếu cần, tôi có thể thực hiện A=B=C (qua số lượng nhỏ đệm).

# 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) 

Answer 1

Nếu bạn đang sử dụng NumPy> 1.6, đó là tuyệt vời np.einsum chức năng:

np.einsum('im,jm,km->ijk',x,y,z) 

Đó là tương đương với phiên bản looped của bạn. Tôi không chắc làm thế nào điều này sẽ công bằng về hiệu quả một khi bạn nhận được lên đến kích thước của mảng của bạn trong vấn đề thực sự (Tôi thực sự nhận được một segfault trên máy tính của tôi, khi tôi di chuyển đến những kích cỡ). Các giải pháp khác mà tôi thường thích cho các loại vấn đề là viết lại phương pháp sử dụng cython.

Answer 2

Sử dụng einsum rất có ý nghĩa trong trường hợp của bạn; nhưng bạn có thể làm điều này khá dễ dàng bằng tay. Bí quyết là làm cho các mảng phát sóng được với nhau. Điều đó có nghĩa là định hình lại chúng sao cho mỗi mảng thay đổi độc lập dọc trục riêng của nó. Sau đó nhân chúng với nhau, cho phép numpy chăm sóc phát sóng; và sau đó tổng hợp dọc theo trục cuối cùng (bên phải).

>>> x = numpy.arange(2 * 4).reshape(2, 4) 
>>> y = numpy.arange(3 * 4).reshape(3, 4) 
>>> z = numpy.arange(4 * 4).reshape(4, 4) 
>>> (x.reshape(2, 1, 1, 4) * 
... y.reshape(1, 3, 1, 4) * 
... z.reshape(1, 1, 4, 4)).sum(axis=3) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]]) 

Bạn có thể làm điều này một chút khái quát hóa bằng cách sử dụng ký hiệu lát, các newaxis giá trị (tương đương với None, vì vậy dưới đây sẽ làm việc với None cũng), và thực tế là sum chấp nhận các giá trị trục tiêu cực (với -1 biểu thị kết quả cuối cùng, -2 biểu thị từ tiếp theo đến lần cuối, v.v.). Bằng cách này, bạn không cần phải biết hình dạng ban đầu của mảng; miễn là trục cuối cùng của chúng tương thích, điều này sẽ phát sóng ba đầu tiên cùng nhau:

>>> (x[:, numpy.newaxis, numpy.newaxis, :] * 
... y[numpy.newaxis, :, numpy.newaxis, :] * 
... z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]])