Numpy 如何根据索引创建变量矩阵

我正在尝试使用

scipy或numpy

创建以下类型的矩阵

A_1 = diag(0.5, 0, 0, ...., 0)
A_k = diag(0,0, ..., 1, 0, 0, ....,0) for each 1 < k < N
A_N = diag(0, 0, ...., 0, 0.5)

或

=.5

如果

k==0

或

k==n-1

还是我遗漏了一些细微差别

要使其稀疏，请从大小合适的

lil

开始。尽管为单个对角线值指定

coo

输入也很容易-

def diag_ndarray(N):
    # Initialize o/p array
    out = np.zeros((N,N,N))

    # Get n-dimensional indices for assigning non-zero values
    idx = np.tile(np.arange(N),3).reshape(3,-1)

    # Setup non-zero values and assign with subscripted indexing
    out[tuple(idx)] = np.hstack((0.5,[1]*(N-2),0.5))
    return out

更简单的版本-

def diag_ndarray_v2(N):
    out = np.zeros((N,N,N))
    r = np.arange(N)
    out[r,r,r] = np.hstack((0.5,[1]*(N-2),0.5))
    return out

更简短的版本，带有线性索引-

def diag_ndarray_v3(N):
    out = np.zeros((N,N,N))
    out.ravel()[np.arange(N)*(1+N*(N+1))] = np.hstack((0.5,[1]*(N-2),0.5))
    return out

样本运行-

In [128]: diag_ndarray(N=4)
Out[128]: 
array([[[ 0.5,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  1. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  1. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0.5]]])

---

不是每个

A_1

等都是二维数组吗？也许给我们看看N=4或什么的预期o/p？是的。我明确地添加了矩阵。为了使其稀疏，我可以从前面建议的适当大小的

lil

开始吗？@JackDawkins认为你不能用任何稀疏矩阵表示，因为它们都不支持3D数组，这是这里所需的输出。这不会创建一系列矩阵，例如，我可以在右边求和？。（例如，sum（A_k）将给我A_1+A_2+A_3+…等等。我如何扩展它？

def diag_ndarray_v3(N):
    out = np.zeros((N,N,N))
    out.ravel()[np.arange(N)*(1+N*(N+1))] = np.hstack((0.5,[1]*(N-2),0.5))
    return out

In [128]: diag_ndarray(N=4)
Out[128]: 
array([[[ 0.5,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  1. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  1. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ]],

       [[ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0. ],
        [ 0. ,  0. ,  0. ,  0.5]]])