涉及numpy数组索引的干净求和

时间:2016-01-19 11:21:09

标签: python arrays numpy

我偶尔也不经常使用numpy。我现在需要做一些总结,其中总和涉及行/列索引。

我有一个m x n数组S.我想创建一个新的m x n数组,其's,i'条目是

-c i S[s,i] + g (i+1)S[s,i+1] + (s+1)S[s+1,i-1]

所以说S=np.array([[1,2],[3,4], [5,6]])我想要的结果是

-c*np.array([[0*1, 1*2],[0*3, 1*4],[0*5, 1*6]]) 
+ g*np.array([[1*2, 2*0],[1*4, 2*0],[1*6, 2*0]]) 
+ np.array([[1*0, 1*3],[2*0, 2*5],[3*0, 3*0]])

(这不是我方程中的所有术语,但我觉得知道如何做到这一点就足以完成我所追求的目标了。)

我认为我需要做的是创建一个新数组,其行只是行的索引,另一行是列的对应。然后做一些分量乘法。但这远远超出了我在研究中通常做的事情,所以我已经采取了一些错误的步骤。

note :据了解,索引引用了我的数组之外的东西,其值为零。

有没有一种干净的方法来完成我上面描述的总结?

1 个答案:

答案 0 :(得分:1)

由于您可能的越界索引,我会分几步完成:

import numpy as np
S = np.array([[1,2],[3,4], [5,6]])
c = np.random.rand()
g = np.random.rand()

m,n = S.shape

Stmp1 = S*np.arange(0,n)         # i*S[s,i]
Stmp2 = S*np.arange(0,m)[:,None] # s*S[s,i]

# the answer:
Sout = -c*Stmp1
Sout[:,:-1] = Sout[:,:-1] + g*Stmp1[:,1:]
Sout[:-1,1:] = Sout[:-1,1:] + Stmp2[1:,:-1]

# only for control:
Sout2 = -c*np.array([[0*1, 1*2],[0*3, 1*4],[0*5, 1*6]]) \
       + g*np.array([[1*2, 2*0],[1*4, 2*0],[1*6, 2*0]]) \
       +   np.array([[1*0, 1*3],[2*0, 2*5],[3*0, 3*0]])

检查:

In [431]: np.all(Sout==Sout2)
Out[431]: True

我为i*S[s,i]s*S[s,i]引入了辅助数组。虽然这显然不是必需的,但它使代码更容易阅读。我们可以直接轻松切入np.arange(0,n)调用,但除非内存不是问题,否则我觉得这种方法更直接。