Как векторизовать несколько уровней рекурсии?

Question

Я noobie для python и numpy (и программирования вообще). Я стараюсь максимально ускорить мой код. Математика включает в себя несколько суммирования по нескольким осям нескольких массивов. Я достиг одного уровня векторизации, но мне кажется, что я не могу глубже понять это и должен прибегать к циклам (я считаю, что есть три уровня рекурсии: M, N и я, один из которых я исключено, I). Вот мой код для соответствующего раздела (этот код работает, но я хотел бы, чтобы ускорить его):

def B1(n, i): 
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1) 

for n in N: 
    B[n, :] = B1(n, I) 

for m in M: 
    for n in N: 
     C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :]) 

    Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :]) 

A = np.linalg.solve(C[1:, 1:], (0.25)*Y[1:]) 

dmaxi только поплавок и т, п и я целые числа. Массивы имеют следующие формы:

>>> qi.shape 
(551,) 
>>> N.shape 
(18,) 
>>> M.shape 
(18,) 
>>> I.shape 
(551,) 
>>> Iq.shape 
(1, 551, 3) 
>>> U.shape 
(1, 551, 3) 

Как вы можете видеть, я векторизованы расчет 2-ой оси В, но я не могу сделать это за 1-ой оси, С и Y, которые по-прежнему требуют циклов for. Похоже, что когда я пытаюсь сделать ту же самую форму векторизации, которую я сделал для 1-й оси B (определите функцию, затем передайте массив как аргумент), я получаю ошибку вещания, так как она пытается вычислить как осей одновременно, а не 1-го, затем 2-го, поэтому мне пришлось вынудить его в цикл for. Такая же проблема возникает как для C, так и для Y, поэтому они также подходят для циклов. В случае это сбивает с толку, по существу, что я попытался было:

>>> B[:, :] = B1(N, I) 
Traceback (most recent call last): 
    File "<stdin>", line 1, in <module> 
    File "sasrec_v6.py", line 155, in B1 
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1) 
ValueError: operands could not be broadcast together with shapes (18) (551) 

Векторизация 2-ой оси B сделал значительное улучшение в скорости моего кода, поэтому я предполагаю, что то же самое будет применяться для дальнейшего векторизации (I надеюсь, что я правильно использую этот термин).

Answer 1

Вы можете использовать broadcasting, чтобы сделать 2d массивы из ваших индексов индекса 1d. Я не проверял это еще, но они должны работать:

Если вы перекроить N быть вектор-столбец, то B1 будет возвращать 2d массив:

B[N] = B1(N[:, None], I) 

Для Y и C, я d использовать np.einsum, чтобы лучше контролировать какую ось mulitplied (вероятно, это может быть сделано с np.dot как хорошо, но я не знаю, как.

C[M[:, None], N] = np.einsum('ij,kj->ik', 
     B[M]/np.square(qi*Iq[0, :, 2]), 
     B[N]) 

Y[M] = np.einsum('i, ki->k', 
     U[0, :, 1]/np.square(qi*Iq[0, :, 2]), 
     B[M]) 

чтобы посмотреть, что индекс ИНГ трюк делает:

In [1]: a = np.arange(3) 

In [2]: a 
Out[2]: array([0, 1, 2]) 

In [3]: a[:, None] 
Out[3]: 
array([[0], 
     [1], 
     [2]]) 

In [4]: b = np.arange(4,1,-1) 

In [5]: b 
Out[5]: array([4, 3, 2]) 

In [6]: a[:, None] * b 
Out[6]: 
array([[0, 0, 0], 
     [4, 3, 2], 
     [8, 6, 4]]) 

Это экономит два порядка во времени:

In [92]: %%timeit 
    ....: B = np.zeros((18, 551)) 
    ....: C = np.zeros((18, 18)) 
    ....: Y = np.zeros((18)) 
    ....: for n in N: 
    ....:  B[n, :] = B1(n, I) 
    ....: for m in M: 
    ....:  for n in N: 
    ....:   C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :]) 
    ....:  Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :]) 
    ....: 
100 loops, best of 3: 15.8 ms per loop 

In [93]: %%timeit 
    ....: Bv = np.zeros((18, 551)) 
    ....: Cv = np.zeros((18, 18)) 
    ....: Yv = np.zeros((18)) 
    ....: Bv[N] = B1(N[:, None], I) 
    ....: Cv[M[:, None], N] = np.einsum('ij,kj->ik', B[M]/np.square(qi*Iq[0, :, 2]), B[N]) 
    ....: Yv[M] = np.einsum('i, ki->k', U[0, :, 1]/np.square(qi*Iq[0, :, 2]), B[M]) 
    ....: 
1000 loops, best of 3: 1.34 ms per loop 

Вот мой тест:

import numpy as np 

# make fake data: 
np.random.seed(5) 

qi = np.random.rand(551) 
N = np.random.randint(0,18,18)#np.arange(18) 
M = np.random.randint(0,18,18)#np.arange(18) 
I = np.arange(551) 
Iq = np.random.rand(1, 551, 3) 
U = np.random.rand(1, 551, 3) 

B = np.zeros((18, 551)) 
C = np.zeros((18, 18)) 
Y = np.zeros((18)) 
Bv = np.zeros((18, 551)) 
Cv = np.zeros((18, 18)) 
Yv = np.zeros((18)) 

dmaxi = 1. 

def B1(n, i): 
    return np.pi * n * dmaxi * (-1)**(n+1) * np.sin(qi[i]*dmaxi) * ((np.pi*n)**2 - (qi[i]*dmaxi)**2)**(-1) 

for n in N: 
    B[n, :] = B1(n, I) 

for m in M: 
    for n in N: 
     C[m, n] = np.dot((1/np.square(qi*Iq[0, :, 2]))*B[m, :], B[n, :]) 
    Y[m] = np.dot((1/np.square(qi*Iq[0, :, 2]))*U[0, :, 1], B[m, :]) 

Bv[N] = B1(N[:, None], I) 
print "B correct?", np.allclose(Bv, B) 

# np.einsum test case: 
n, m = 2, 3 
a = np.arange(n*m).reshape(n,m)*8 + 2 
b = np.arange(n*m)[::-1].reshape(n,m) 
c = np.empty((n,n)) 
for i in range(n): 
    for j in range(n): 
     c[i,j] = np.dot(a[i],b[j]) 
cv = np.einsum('ij,kj->ik', a, b) 
print "einsum test successful?", np.allclose(c,cv) 

Cv[M[:, None], N] = np.einsum('ij,kj->ik', 
     B[M]/np.square(qi*Iq[0, :, 2]), 
     B[N]) 
print "C correct?", np.allclose(Cv, C) 

Yv[M] = np.einsum('i, ki->k', 
     U[0, :, 1]/np.square(qi*Iq[0, :, 2]), 
     B[M]) 
print "Y correct?", np.allclose(Yv, Y) 

выход: D

B correct? True 
einsum test successful? True 
C correct? True 
Y correct? True