在Julia中制作n维对角张量

Question

例如，让我们用ijk表示3d张量的第a[i][j][k]个元素。可以说3d张量在每个维度上的大小都相同。然后，我想用a[i][i][i]用constant[i]填充i，从1到n。我希望所有其他元素都为0。唯一的方法是使for循环吗？

Answer 1

您可以理解：

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

Answer 2

如果a已经存在（例如a = zeros(n, n, n)），则可以使用以下函数索引主对角线，该主对角线应适用于任何尺寸（也适用于非对称情况）：< / p>

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

然后您可以简单地广播：

a[diagind(a)] .= constant

如果始终可以将输入假定为对角线，那么minimum(size(A))当然可以仅用N代替。