在numpy中構造這個數組的有效方法？

Question

我需要構造一個N x M x N陣列A ，使得A[i, j, k] == (0 if i != k else x[j]) 。 我可以寫：

A = np.zeros((N, M, N))
for i in range(N):
    for j in range(M):
        A[i,j,i] = x[j]

或者，或者：

A = np.zeros((N, M, N))
for i in range(N):
    A[i,:,i] = x

但兩者都很可能對我的目的來說太慢了。 有更快的方法嗎？

Answer 1

方法＃1

使用broadcasting來創建所有那些要分配x的線性索引，然后簡單地分配到它的展平視圖中，就像這樣 -

# Initialize
Aout = np.zeros((N, M, N))

# Comput all linear indices
idx = (np.arange(N)*(N*M+1))[:,None] + N*np.arange(M)

# In a flattened view with `.ravel()` assign from x
Aout.ravel()[idx] = x

方法＃2

使用np.lib.stride_tricks.as_strided支持的基於視圖的元素訪問 -

Aout = np.zeros((N, M, N))
s0,s1,s2 = Aout.strides
Aoutview = np.lib.stride_tricks.as_strided(Aout,shape=(N,M),strides=(s0+s2,s1))
Aoutview[:] = x

方法＃3

另一種方法是沿第一和第三軸使用integer array indexing ，從問題中模擬第二種方法，但是以矢量化方式 -

Aout = np.zeros((N, M, N))
Aout[np.arange(N),:,np.arange(N)] = x

運行時測試

方法 -

def app0(x,A):
    for i in range(N):
        for j in range(M):
            A[i,j,i] = x[j]
    return A

def app1(x,A):
    for i in range(N):
        A[i,:,i] = x
    return A

def app2(x,Aout):       
    idx = (np.arange(N)*(N*M+1))[:,None] + N*np.arange(M)
    Aout.ravel()[idx] = x
    return Aout

def app3(x,Aout):       
    s0,s1,s2 = Aout.strides
    Aoutview = np.lib.index_tricks.as_strided(Aout,shape=(N,M),strides=(s0+s2,s1))
    Aoutview[:] = x
    return Aout

def app4(x,Aout):
    r = np.arange(N)
    Aout[r,:,r] = x
    return Aout

驗證 -

In [125]: # Params
     ...: N, M = 100,100
     ...: x = np.random.rand(M)
     ...: 
     ...: # Make copies of arrays to be assigned into
     ...: A0 = np.zeros((N, M, N))
     ...: A1 = np.zeros((N, M, N))
     ...: A2 = np.zeros((N, M, N))
     ...: A3 = np.zeros((N, M, N))
     ...: A4 = np.zeros((N, M, N))
     ...: 

In [126]: print np.allclose(app0(x,A0), app1(x,A1))
     ...: print np.allclose(app0(x,A0), app2(x,A2))
     ...: print np.allclose(app0(x,A0), app3(x,A3))
     ...: print np.allclose(app0(x,A0), app4(x,A4))
     ...: 
True
True
True
True

計時 -

In [127]: # Make copies of arrays to be assigned into
     ...: A0 = np.zeros((N, M, N))
     ...: A1 = np.zeros((N, M, N))
     ...: A2 = np.zeros((N, M, N))
     ...: A3 = np.zeros((N, M, N))
     ...: A4 = np.zeros((N, M, N))


In [128]: %timeit app0(x,A0)
     ...: %timeit app1(x,A1)
     ...: %timeit app2(x,A2)
     ...: %timeit app3(x,A3)
     ...: %timeit app4(x,A4)
     ...: 
1000 loops, best of 3: 1.49 ms per loop
10000 loops, best of 3: 53.6 µs per loop
10000 loops, best of 3: 150 µs per loop
10000 loops, best of 3: 28.6 µs per loop
10000 loops, best of 3: 25.2 µs per loop