Сделайте n размерный диагональный тензор в Юлии

Question

Например, давайте обозначим ijk '-й элемент трехмерного тензора a[i][j][k].Допустим, 3d-тензор имеет одинаковый размер в каждом измерении.Затем я хочу заполнить a[i][i][i] на constant[i] для i от 1 до n.Я хочу, чтобы все остальные элементы были равны 0. Является ли создание цикла for моим единственным способом?

Answer 1

Если a уже существует (например, a = zeros(n, n, n)), вы можете использовать следующую функцию для индексации главной диагонали, которая должна работать для любого измерения (а также для несимметричных случаев):

diagind(A::AbstractArray{T,N} where T) where N = CartesianIndex.(fill(1:minimum(size(A)), N)...)

Тогда вы можете просто передать:

a[diagind(a)] .= constant

В случае, если вход всегда можно считать диагональным, minimum(size(A)), конечно, можно заменить просто N.

Answer 2

Вы можете использовать понимание:

a = [i==j==k ? constant[i] : 0.0 for i=1:n, j=1:n, k=1:n]