Я пробовал что-то вроде
ret = np.zeros((M,N,M,N))
ind = np.arange(M)
ret[ind,:,ind,:] = tempH0s[ind,:,:]
ret = ret.reshape(M*N, M*N)
, но расположение значений отличалось от того, что делает ваша итерация. Я подозреваю, что отображение in1+M*in2
является частью разницы.
Если разницу можно устранить, это должно быть немного быстрее.
In [93]: tempH0s
Out[93]:
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]]])
ваша итерация дает
In [87]: foo(tempH0s,2,3)
Out[87]:
array([[ 0, 1, 0, 0, 0, 0],
[ 3, 4, 0, 0, 0, 0],
[ 0, 0, 9, 10, 11, 0],
[ 0, 0, 12, 13, 14, 0],
[ 0, 0, 15, 16, 17, 0],
[ 0, 0, 0, 0, 0, 0]])
, где, как я предлагаю, дает:
array([[ 0, 1, 2, 0, 0, 0],
[ 3, 4, 5, 0, 0, 0],
[ 6, 7, 8, 0, 0, 0],
[ 0, 0, 0, 9, 10, 11],
[ 0, 0, 0, 12, 13, 14],
[ 0, 0, 0, 15, 16, 17]])
Ваш другой код дает
array([[ 0., 0., 1., 0., 2., 0.],
[ 0., 9., 0., 10., 0., 11.],
[ 3., 0., 4., 0., 5., 0.],
[ 0., 12., 0., 13., 0., 14.],
[ 6., 0., 7., 0., 8., 0.],
[ 0., 15., 0., 16., 0., 17.]])