Обработка верхних треугольных элементов только с помощью NumPy einsum

Question

Я использую numy einsum для вычисления точечных произведений массива векторов столбцов pts, формы (3,N), с самим собой, в результате чего получается матричная точка, формы (N,N) со всеми точечными произведениями. Это код, который я использую:

dotps = np.einsum('ij,ik->jk', pts, pts)

Это работает, но мне нужны только значения выше главной диагонали. то есть. верхняя треугольная часть результата без диагонали. Можно ли вычислить только эти значения с помощью einsum? или любым другим способом, который быстрее, чем использование einsum для вычисления всей матрицы?

Мой массив pts может быть довольно большим, поэтому, если бы я мог рассчитывать только те значения, которые мне нужны, это удвоило бы мою скорость вычислений.

Answer 1

Вы можете нарезать соответствующие столбцы, а затем использовать np.einsum -

R,C = np.triu_indices(N,1)
out = np.einsum('ij,ij->j',pts[:,R],pts[:,C])

Пробный прогон -

In [109]: N = 5
     ...: pts = np.random.rand(3,N)
     ...: dotps = np.einsum('ij,ik->jk', pts, pts)
     ...: 

In [110]: dotps
Out[110]: 
array([[ 0.26529103,  0.30626052,  0.18373867,  0.13602931,  0.51162729],
       [ 0.30626052,  0.56132272,  0.5938057 ,  0.28750708,  0.9876753 ],
       [ 0.18373867,  0.5938057 ,  0.84699103,  0.35788749,  1.04483158],
       [ 0.13602931,  0.28750708,  0.35788749,  0.18274288,  0.4612556 ],
       [ 0.51162729,  0.9876753 ,  1.04483158,  0.4612556 ,  1.82723949]])

In [111]: R,C = np.triu_indices(N,1)
     ...: out = np.einsum('ij,ij->j',pts[:,R],pts[:,C])
     ...: 

In [112]: out
Out[112]: 
array([ 0.30626052,  0.18373867,  0.13602931,  0.51162729,  0.5938057 ,
        0.28750708,  0.9876753 ,  0.35788749,  1.04483158,  0.4612556 ])

---

Оптимизация дальше -

Давайте рассмотрим наш подход и посмотрим, есть ли возможности для улучшения производительности.

In [126]: N = 5000

In [127]: pts = np.random.rand(3,N)

In [128]: %timeit np.triu_indices(N,1)
1 loops, best of 3: 413 ms per loop

In [129]: R,C = np.triu_indices(N,1)

In [130]: %timeit np.einsum('ij,ij->j',pts[:,R],pts[:,C])
1 loops, best of 3: 1.47 s per loop

Оставаясь в рамках ограничений памяти, похоже, что мы ничего не можем сделать с оптимизацией np.einsum, Итак, давайте сместим акцент на np.triu_indices,

За N = 4, у нас есть:

In [131]: N = 4

In [132]: np.triu_indices(N,1)
Out[132]: (array([0, 0, 0, 1, 1, 2]), array([1, 2, 3, 2, 3, 3]))

Кажется, что он создает регулярную модель, вроде как изменяющуюся. Это может быть записано с накопленной суммой, которая имеет сдвиги в тех 3 а также 5 позиции. Думая в общем, мы бы закончили кодировать что-то вроде этого -

def triu_indices_cumsum(N):

    # Length of R and C index arrays
    L = (N*(N-1))/2

    # Positions along the R and C arrays that indicate 
    # shifting to the next row of the full array
    shifts_idx = np.arange(2,N)[::-1].cumsum()

    # Initialize "shift" arrays for finally leading to R and C
    shifts1_arr = np.zeros(L,dtype=int)
    shifts2_arr = np.ones(L,dtype=int)

    # At shift positions along the shifts array set appropriate values, 
    # such that when cumulative summed would lead to desired R and C arrays.
    shifts1_arr[shifts_idx] = 1
    shifts2_arr[shifts_idx] = -np.arange(N-2)[::-1]

    # Finall cumsum to give R, C
    R_arr = shifts1_arr.cumsum()
    C_arr = shifts2_arr.cumsum()
    return R_arr, C_arr

Давайте время для различных N's!

In [133]: N = 100

In [134]: %timeit np.triu_indices(N,1)
10000 loops, best of 3: 122 µs per loop

In [135]: %timeit triu_indices_cumsum(N)
10000 loops, best of 3: 61.7 µs per loop

In [136]: N = 1000

In [137]: %timeit np.triu_indices(N,1)
100 loops, best of 3: 17 ms per loop

In [138]: %timeit triu_indices_cumsum(N)
100 loops, best of 3: 16.3 ms per loop

Таким образом, это выглядит как приличный N's, индивидуальные cumsum на основе triu_indices может быть стоит посмотреть!