Вычислить след матрицы по всем диагоналям

Если вы хотите держаться подальше от Cython, построив диагональный массив индексов и используя np.bincount может сделать трюк:

>>> import numpy as np
>>> a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> rows, cols = a.shape
>>> rows_arr = np.arange(rows)
>>> cols_arr = np.arange(cols)
>>> diag_idx = rows_arr[:, None] - (cols_arr - (cols - 1))
>>> diag_idx
array([[3, 2, 1, 0],
       [4, 3, 2, 1],
       [5, 4, 3, 2]])
>>> np.bincount(diag_idx.ravel(), weights=a.ravel())
array([  3.,   9.,  18.,  15.,  13.,   8.])

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

Если ваша форма матрицы достаточно далека от квадрата, то есть если она высокая или широкая, то вы можете эффективно использовать трюки stride для этого. Вы можете использовать трюки stride в любом случае, но это может быть не очень эффективной памятью, если матрица находится рядом с квадратом.

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

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

Если вам не нужны меньшие следы по краям, то вам не нужно нулевую площадку.

здесь идет (если больше столбцов, чем строк, но легко адаптирован):

import numpy as np
from numpy.lib.stride_tricks import as_strided

A = np.arange(30).reshape(3, 10)
A_embedded = np.hstack([np.zeros([3, 2]), A, np.zeros([3, 2])])
A = A_embedded[:, 2:-2]  # We are now sure that the memory around A is padded with 0, but actually we never really need A again

new_strides = (A.strides[0] + A.strides[1], A.strides[1])
B = as_strided(A_embedded, shape=A_embedded[:, :-2].shape, strides=new_strides)

traces = B.sum(0)

print A
print B
print traces

для того, чтобы соответствуйте выходу, который вы показываете в своем примере, вам нужно отменить его (см. комментарий @larsmans)

traces = traces[::-1]

Это конкретный пример с конкретными цифрами. Если это полезно для вашего использования, я могу превратить его в общую функцию.

это конкурентоспособно, если массив большой:

def f5(A):
    rows, cols = A.shape
    N = rows + cols -1
    out = np.zeros(N, A.dtype)
    for idx in range(rows):
        out[N-idx-cols:N-idx] += A[idx]
    return out[::-1]

хотя он использует цикл Python, он быстрее, чем bincount решение (для больших массивов.. на моей системе..)

этот метод имеет высокую чувствительность к соотношению столбец/строка массива, потому что это соотношение определяет, сколько цикла выполняется в Python относительно Numpy. Как отметил @Jaime, эффективно перебирать наименьшее измерение, например:

def f6(A):
    rows, cols = A.shape
    N = rows + cols -1
    out = np.zeros(N, A.dtype)

    if rows > cols:
        for idx in range(cols):
            out[N-idx-rows:N-idx] += A[:, idx]
    else:
        for idx in range(rows):
            out[N-idx-cols:N-idx] += A[idx]
        out = out[::-1]
    return out

но следует отметить, это для больших размеров массива (например,100000 x 500 в моей системе) доступ к массиву строка за строкой, как в первом коде, который я опубликовал, все еще может быть быстрее, вероятно, из-за того, как массив выложен в ОЗУ (быстрее получать непрерывные куски, чем разбросанные биты).

Это можно сделать (слегка оскорбительно) с помощью scipy.sparse.dia_matrix двумя способами, один реже, чем другой.

первый, давая точный результат, использует dia_matrix вектор сохраненных данных

import numpy as np
from scipy.sparse import dia_matrix
A = np.arange(30).reshape(3, 10)
traces = dia_matrix(A).data.sum(1)[::-1]

менее интенсивным методом памяти было бы работать наоборот:

import numpy as np
from scipy.sparse import dia_matrix
A = np.arange(30).reshape(3, 10)
A_dia = dia_matrix((A, range(len(A))), shape=(A.shape[1],) * 2)
traces = np.array(A_dia.sum(1)).ravel()[::-1]

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

@moarningsun найдено решение:

rows, cols = A.shape

A_dia = dia_matrix((A, np.arange(rows)), shape=(cols,)*2)
traces1 = A_dia.sum(1).A.ravel()

A_dia = dia_matrix((A, np.arange(-rows+1, 1)), shape=(rows,)*2)
traces2 = A_dia.sum(1).A.ravel()

traces = np.concatenate((traces1[::-1], traces2[-2::-1]))

np.trace делает то, что вы хотите:

import numpy as np

A = array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]])

n = A.shape[0]
[np.trace(A, i) for i in range(-n+1, n+1)]

редактировать: изменить np.sum(np.diag()) to np.trace() согласно предложению от @user2357112.

используйте массив numpy trace способ:

import numpy as np
A = np.array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]])
A.trace()

возвращает:

15