2015-04-06 54 views
1

我已閱讀此post,並懷疑我需要更好地理解索引。我想這樣做np.ndenumerate,但(在這種特殊情況下)僅在第一三個維度,返回了一系列載體:像部分/選擇np.ndenumerate?

x = np.linspace(0., 1., 4) 
y = np.linspace(0., 1., 5) 
z = np.linspace(0., 1., 2) 
X, Y, Z = np.meshgrid(x, y, z) # define a grid 
F = np.zeros(X.shape + (3,))  # a 3D vector field 
F = np.random.rand(5*4*2*3).reshape(5,4,2,3) # added this later just to give non-zero for testing 

thetas = np.linspace(0, 2.*np.pi, 21)[:-1] # drop the duplicate 
v = np.array([(np.cos(theta), np.sin(theta), 0.0) for theta in thetas]) 

for tup, vec in magical_enumerate(F, axis=(0,1,2)): # it's magic! (need correct expression) 
    F(tup) = np.cross(v, vec).sum(axis=0)  # sum over 20 vectors in v 

有沒有辦法做到這一點沒有大量的循環或列表的解釋?網格將很大,所以numpythony和速度值得讚賞。非連續維度(例如,軸=(0,2))呢?謝謝!

+0

你可以發佈可運行代碼,因爲'magical_enumerate'是未定義的嗎? – Divakar 2015-04-06 06:18:02

+0

嗨Divarak,這是我在SO中看到的那種速記,因爲「我不知道該把什麼放在這裏,但這是我正在尋找的東西。」該代碼僅僅是一個幫助澄清文本的例子。 – uhoh 2015-04-06 07:43:19

+0

我可能已經弄錯了,但是你不能只是重塑'F',然後像'np.sum(np.cross(F.reshape(-1,F.shape [-1 ]),vec),1)'然後用這些值設置爲'F'? – Divakar 2015-04-06 09:18:51

回答

3

np.ndindex可能會訣竅。它通過一組維度生成一個迭代器。

In [231]: F=np.zeros((2,3,3,3)) 

In [232]: for tup in np.ndindex(F.shape[:3]): 
    # vec = data[tup] etc 
    F[tup]=tup 
    .....:  

In [233]: F 
Out[233]: 
array([[[[ 0., 0., 0.], 
     [ 0., 0., 1.], 
     [ 0., 0., 2.]], 

     ... 
     [ 1., 2., 1.], 
     [ 1., 2., 2.]]]]) 

我建議看代碼都ndenumeratendindexndindexmulti_index模式下使用最近的nditerndenumerate使用flat來遍歷所有值。

我建議在其他所以你可以如何構建自己的multi_index迭代器模仿ndindex。在nditer上搜索可能會產生這些結果。

這不會給多個循環帶來速度優勢,因爲您仍然在處理相同數量的最內部計算。

至於非順序的尺寸,這同樣ndindex將工作,但你必須使用它作爲索引之前操縱tup

In [243]: for tup in np.ndindex((F.shape[0],F.shape[2])): 
    tup1=(tup[0],slice(None),tup[1]) 
    F[tup]=[tup[0],np.nan,tup[1]] 
    .....:  

np.apply_along_axisnp.apply_over_axes超過1個或多個生成的索引的其他例子在切割其他切片的同時,

+0

謝謝hpaulj!確實很有效!現在,如果我可以弄清楚如何用'[:]'而不是循環來做到這一點,速度會更快! – uhoh 2015-04-06 08:23:21