Создать стек массивов из диагональных значений, используя NumPy

Question

Я пытаюсь выполнить некоторые матричные вычисления в python и столкнулся с проблемой, когда попытался ускорить мой код, используя составные массивы вместо простых циклов for.Мне нужно создать 2D-массив со значениями (заданными в виде 1D-массива) по диагонали, но я не смог придумать, как сделать это со сложенными массивами.

В старой (циклической) версииЯ использовал метод np.diag (), который возвращает именно то, что мне нужно (2D-массив в этом случае), если я задаю значения как 1D-массив в качестве входных данных.Однако, когда я переключился на составные массивы, мой ввод больше не является 1D-массивом, так что метод np.diag () возвращает копию диагонали моего 2D-ввода.

Старая версия с 1Dinput:

import numpy as np
vals = np.array([1,2,3])
mat = np.diag(vals)
print(mat.shape)

Out: (3, 3)

Новая версия с 2D-вводом:

vals_stack = np.repeat(np.expand_dims(vals, axis=0), 5, axis=0) 
# btw: is there a better way to repeat/stack my array? 
mat_stack = np.diag(vals_stack)
print(mat_stack.shape)

Out: (3,)

Итак, вы можете видеть, что np.diag () возвращает 1D-массив (как и ожидалось)из документации), но мне на самом деле нужен стек 2D-массивов.Таким образом, форма mat_stack должна быть (7,3,3), а не (3,).Есть ли какая-нибудь функция для этого в NumPy?Или я должен зациклить это дополнительное измерение следующим образом:

def mydiag(stack):
    diag = np.zeros([stack.shape[0], stack.shape[1], stack.shape[1]])
    for i in np.arange(stack.shape[0]):
        diag[i,:,:] = np.diag([stack[i,:].ravel()])
    return diag

Answer 1

В numpy вы должны использовать apply_along_axis.В конце документа есть даже пример для вашего конкретного случая ( здесь ).Таким образом, ответ:

np.apply_along_axis(np.diag, -1, vals_stack)

Более питонический путь будет примерно таким:

[np.diag(row) for row in vals_stack]

Answer 2

Это то, что вы имели в виду:

In [499]: x = np.arange(12).reshape(4,3)                                                                     
In [500]: X = np.zeros((4,3,3),int)                                                                          
In [501]: X[np.arange(4)[:,None],np.arange(3), np.arange(3)] = x                                             
In [502]: X                                                                                                  
Out[502]: 
array([[[ 0,  0,  0],
        [ 0,  1,  0],
        [ 0,  0,  2]],

       [[ 3,  0,  0],
        [ 0,  4,  0],
        [ 0,  0,  5]],

       [[ 6,  0,  0],
        [ 0,  7,  0],
        [ 0,  0,  8]],

       [[ 9,  0,  0],
        [ 0, 10,  0],
        [ 0,  0, 11]]])

X[0,np.arange(3), np.arange(3)] индексирует диагональ на первой плоскости. np.arange(4)[:,None] - это массив (4,1), который вещает с (3,) для индексации блока (4,3), соответствующего размеру x.