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

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

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

Старая версия с 1D вводом:

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

python numpy numpy-ndarray
person koxx    schedule 09.07.2019    source источник

Ответы (2)


arrow_upward
1
arrow_downward

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

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

Более питоническим способом было бы что-то вроде этого:

[np.diag(row) for row in vals_stack]
person Simon Delecourt    schedule 09.07.2019
comment
Это именно то, что я искал, и это работает как шарм! Благодарю вас! - person koxx; 10.07.2019

arrow_upward
0
arrow_downward

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

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.

person hpaulj    schedule 09.07.2019