ファインダーのインデックス作成と縮小は高速ですか？

Question

私は、2つの配列を「結合」し、結果の軸の1つを合計するために、ファンシーインデックスを使用して加速しようとしています。このような

何か：

$ ipython 
In [1]: import numpy as np 
In [2]: ne, ds = 12, 6 
In [3]: i = np.random.randn(ne, ds).astype('float32') 
In [4]: t = np.random.randint(0, ds, size=(1e5, ne)).astype('uint8') 

In [5]: %timeit i[np.arange(ne), t].sum(-1) 
10 loops, best of 3: 44 ms per loop 

はIn [5]で声明を加速するための簡単な方法はありますか？私はOpenMPとか、scipy.weaveやCythonのprangeのようなものを使うべきですか？

Answer 1

は何らかの理由で空想的なインデックス作成よりもはるかに高速です。唯一のトリックは、配列をフラットとして扱うことです。

In [1]: a = np.random.randn(12,6).astype(np.float32) 

In [2]: c = np.random.randint(0,6,size=(1e5,12)).astype(np.uint8) 

In [3]: r = np.arange(12) 

In [4]: %timeit a[r,c].sum(-1) 
10 loops, best of 3: 46.7 ms per loop 

In [5]: rr, cc = np.broadcast_arrays(r,c) 

In [6]: flat_index = rr*a.shape[1] + cc 

In [7]: %timeit a.take(flat_index).sum(-1) 
100 loops, best of 3: 5.5 ms per loop 

In [8]: (a.take(flat_index).sum(-1) == a[r,c].sum(-1)).all() 
Out[8]: True 

私はあなたがこれ以上の速度向上の多くを見に行くしている唯一の他の方法はPyCUDAのようなものを使用したGPU用のカスタムカーネルを書くことになると思います。