在Python中改進for循環的性能（可能使用numpy或numba）

Question

我想提高此函數中for循環的性能。

import numpy as np 
import random 

def play_game(row, n=1000000): 
    """Play the game! This game is a kind of random walk. 

    Arguments: 
     row (int[]): row index to use in the p matrix for each step in the 
        walk. Then length of this array is the same as n. 

     n (int): number of steps in the random walk 
    """ 
    p = np.array([[ 0.499, 0.499, 0.499], 
        [ 0.099, 0.749, 0.749]]) 
    X0 = 100 
    Y0 = X0 % 3 
    X = np.zeros(n) 
    tempX = X0 
    Y = Y0 

    for j in range(n): 
     tempX = X[j] = tempX + 2 * (random.random() < p.item(row.item(j), Y)) - 1 
     Y = tempX % 3 

    return np.r_[X0, X] 

困難在於以下事實的Y值被在每個步驟中，基於的X和的值，該值Y然後在下一步驟中使用，以更新X的值計算的。

我不知道是否有一些可能會造成很大差異的疙瘩技巧。使用Numba是公平的遊戲（我嘗試過但沒有太多成功）。但是，我不想使用Cython。

Answer 1

快速觀察告訴我們，函數代碼中迭代之間存在數據依賴關係。現在，有不同種類的數據依賴關係。您正在查看的數據依賴性的類型是索引依賴性，即在任何迭代中的數據選擇取決於先前的迭代計算。這種依賴似乎很難在迭代之間進行追蹤，所以本文不是真正的矢量化解決方案。相反，我們會盡量預先計算將在循環中使用的值。基本的想法是在循環中做最低限度的工作。

下面是我們如何能夠帶預計算進行，因而具有更有效的解決方案的簡要說明：

• 考慮的p相對較小的形狀從哪一行元素被提取基於輸入row，您可以預先選擇p與p[row]中的所有行。

• 對於每次迭代，您正在計算一個隨機數。你可以用一個隨機數組替換它，你可以在循環之前設置它，因此，你也可以預先計算這些隨機數。

• 根據到目前爲止的預先計算的值，您將擁有p中所有行的列索引。請注意，這些列索引將是一個包含所有可能的列索引的大型ndarray，而在我們的代碼中，只會根據每次迭代計算選擇一個列索引。使用每迭代列索引，您可以增加或減少X0以獲取每次迭代輸出。

的實施應該是這樣的 -

randarr = np.random.rand(n) 
p = np.array([[ 0.499, 0.419, 0.639], 
       [ 0.099, 0.749, 0.319]]) 

def play_game_partvect(row,n,randarr,p): 

    X0 = 100 
    Y0 = X0 % 3 

    signvals = 2*(randarr[:,None] < p[row]) - 1 
    col_idx = (signvals + np.arange(3)) % 3 

    Y = Y0 
    currval = X0 
    out = np.empty(n+1) 
    out[0] = X0 
    for j in range(n): 
     currval = currval + signvals[j,Y] 
     out[j+1] = currval 
     Y = col_idx[j,Y] 

    return out 

爲針對原始代碼驗證，你會修改，像這樣的原碼 -

請注意，因爲這代碼會預先計算這些隨機值，所以這已經可以讓您對問題中的代碼加速很多。

運行測試和驗證輸出 -

In [2]: # Inputs 
    ...: n = 1000 
    ...: row = np.random.randint(0,2,(n)) 
    ...: randarr = np.random.rand(n) 
    ...: p = np.array([[ 0.499, 0.419, 0.639], 
    ...:    [ 0.099, 0.749, 0.319]]) 
    ...: 

In [3]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p)) 
Out[3]: True 

In [4]: %timeit play_game(row,n,randarr,p) 
100 loops, best of 3: 11.6 ms per loop 

In [5]: %timeit play_game_partvect(row,n,randarr,p) 
1000 loops, best of 3: 1.51 ms per loop 

In [6]: # Inputs 
    ...: n = 10000 
    ...: row = np.random.randint(0,2,(n)) 
    ...: randarr = np.random.rand(n) 
    ...: p = np.array([[ 0.499, 0.419, 0.639], 
    ...:    [ 0.099, 0.749, 0.319]]) 
    ...: 

In [7]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p)) 
Out[7]: True 

In [8]: %timeit play_game(row,n,randarr,p) 
10 loops, best of 3: 116 ms per loop 

In [9]: %timeit play_game_partvect(row,n,randarr,p) 
100 loops, best of 3: 14.8 ms per loop 

因此，我們看到大約7.5x+的加速，還不錯！