numpy矢量化從x * y數組到x * y * 3

Question

我有一個函數將整數轉換爲長度爲3的元組，直到現在我一直在爲循環和緩存進行嵌套，但它仍然是很慢。函數pre計算值並將它們緩存在列表中，然後用__getitem__做一些數學運算來確定返回哪個值。然後用不同的數字調用數百萬次。

我最近開始使用numpy，並且我想盡可能地加快這個過程。我做的初始優化（劃分整個numpy數組以匹配緩存長度，所以__getitem__可以被忽略）加速了15到3秒的一個例子，並且如果可能的話我想嘗試使用向量化。

唯一的問題是我無法弄清楚如何使用vectorize將w*h數組轉換爲w*h*3，並且想知道是否有人知道如何？我已經考慮添加一個額外的維度，但它似乎更像是一個重塑類型的東西。

下面是一些示例代碼：

>>> height, width = 2, 5 

>>> array = numpy.arange(width * height).reshape(height, width) 
[[0 1 2 3 4] 
[5 6 7 8 9]] 

>>> test_list = [(0, i, 0) for i in range(width * height)] 
>>> numpy.vectorize(test_list.__getitem__)(array) 
(array([[0, 0, 0, 0, 0], 
     [0, 0, 0, 0, 0]]), array([[0, 1, 2, 3, 4], 
     [5, 6, 7, 8, 9]]), array([[0, 0, 0, 0, 0], 
     [0, 0, 0, 0, 0]])) 

#expected = [[(0, 0, 0) (0, 1, 0) (0, 2, 0) (0, 3, 0) (0, 4, 0)] 
#   [(0, 5, 0) (0, 6, 0) (0, 7, 0) (0, 8, 0) (0, 9, 0)]] 

Answer 1

你可以簡單地在這種情況下使用賦值：

>>> import numpy as np 
>>> height, width = 2, 5 
>>> array = np.arange(width * height).reshape(height, width) 
>>> res = np.zeros((height, width, 3), dtype=array.dtype) 
>>> res[:, :, 1].flat = array.ravel() 
>>> res 
array([[[0, 0, 0], 
     [0, 1, 0], 
     [0, 2, 0], 
     [0, 3, 0], 
     [0, 4, 0]], 

     [[0, 5, 0], 
     [0, 6, 0], 
     [0, 7, 0], 
     [0, 8, 0], 
     [0, 9, 0]]]) 

如果您想元組，你可能會更好過使用普通for循環，普通列表：

>>> lst = array.tolist() 
>>> res = [[(0, item, 0) for item in sublst] for sublst in lst] 
>>> res 
[[(0, 0, 0), (0, 1, 0), (0, 2, 0), (0, 3, 0), (0, 4, 0)], 
[(0, 5, 0), (0, 6, 0), (0, 7, 0), (0, 8, 0), (0, 9, 0)]] 

兩種方法都（可能）更快比np.vectorize這只是一個方便的功能，而不是表現。