numpy-张量乘积

时间:2018-07-30 18:00:51

标签: python numpy tensor

我有一个4 x 4矩阵

import numpy as np
c = np.random.rand((4,4))

我想创建一个100 x 4 x 4 x 100张量,以便当第一个索引与最后一个索引相等时,我得到矩阵,否则得到零。

我可以像这样

Z = np.zeros((100, 4, 4, 100))
for i in range(100):
    Z[i, :, :, i] = c

有更好的方法吗?我尝试查看np.tensordot和np.einsum,但无法弄清楚。

谢谢, 萨希尔

1 个答案:

答案 0 :(得分:1)

使用advanced-indexing-

n = 100
Zout = np.zeros((n, 4, 4, n))
I = np.arange(n)
Zout[I,:,:,I] = c

使用eye-masking-

n = 100
mask = np.eye(n, dtype=bool)
Zout = np.zeros((n, 4, 4, n))
Zout.transpose(0,3,1,2)[mask] = c

时间-

In [72]: c = np.random.rand(4,4)

In [73]: %%timeit
    ...: n = 100
    ...: Zout = np.zeros((n, 4, 4, n))
    ...: I = np.arange(n)
    ...: Zout[I,:,:,I] = c
10000 loops, best of 3: 47.5 µs per loop

In [74]: %%timeit
    ...: n = 100
    ...: mask = np.eye(n, dtype=bool)
    ...: Zout = np.zeros((n, 4, 4, n))
    ...: Zout.transpose(0,3,1,2)[mask] = c
10000 loops, best of 3: 73.1 µs per loop