從 Python + Numba 返回不同的輸出

Question

我有一個 Python function，我正在使用 Numba 優化速度。 但是我需要讓 function 返回一個 float64（它將由需要單個浮點輸出的最小化算法調用）。 但是，一旦優化完成，我還需要它返回有關內部 state 的信息。 這由 26 個 Numpy 向量的元組表示。

目前 Numba 只允許一種類型的 output。

我嘗試使用輸入標志將 state 輸出到文件，從而避免返回類型 nut Numba 不喜歡文件 I/O。

我還能如何訪問內部 state 並仍然使用 Numba？

Answer 1

當您嘗試優化速度時，您不應該涉及文件系統。

如果 26 個向量可以組合成一個 2D Numpy 數組，則可以將其作為另一個參數傳遞給 function。 然后你可以在function里面修改它的內容。

如果您收到錯誤提示您無法修改數組，則需要使用以下命令將其聲明為讀/寫

nb.types.Array(dtype=f8, ndim=2, layout="C", readonly=False)

這也可能適用於 Numba UniTuple 而不是 Numpy 數組。

Answer 2

我使用兩個函數解決了這個問題，目標 function 包裹另一個。

fObj (x, *args) -> float64 返回一個浮點數，是給優化器的浮點數

f(x, output_type, *args) -> float 或 np.array 根據 output 類型標志返回不同的 output。 因此我有看起來像這樣的代碼

@njit 
def fObj(x, args*):    
    output_type = 1    
    output = f(x, output_type, *args)    
    return output

@njit 
def f(x, output_type, args*)

   ... code that calculates objective function ...

   if output_type == 1:
     return xopt
   else:
     return np.array(full_state)