如何在 numba function 中初始化 output numpy arrays？

Question

我有一个计算开角值的 function

@njit(cache=True, nopython=True)
def coord_unit_vec(latlon_vec):

    lat_vec = latlon_vec [:, :, 0]/(latlon_vec [:, :, 0]+latlon_vec [:, :, 1])
    lon_vec = latlon_vec [:, :, 1]/(latlon_vec [:, :, 0]+latlon_vec [:, :, 1])

    return lat_vec, lon_vec

@njit(cache=True, nopython=True)
def calc_oa(latlon, oa_skip):
    X_LAT, X_LON = coord_unit_vec(( latlon[:,2 * oa_skip : ] - latlon[:,oa_skip : -oa_skip]))
    Y_LAT, Y_LON = coord_unit_vec(( latlon[:,: -2 * oa_skip] - latlon[:,oa_skip : -oa_skip]))
    OA = np.empty(dtype=np.float64)
    OA = nan2zero(np.arctan2(Y_LAT, Y_LON) - np.arctan2( X_LAT, X_LON ))
    OA [OA>= np.pi] -= 2*np.pi
    OA [OA<= -1*np.pi] += 2*np.pi
    return np.degrees(OA)

但它不断引发打字错误，没有更多信息到 go

TypingError                               
...
    <source elided>
    OA = np.empty(dtype=np.float64)
    OA = nan2zero(np.arctan2(Y_LAT, Y_LON) - np.arctan2( X_LAT, X_LON ))

请帮助我了解如何更好地打字？

Answer 1

简短回答：您不需要，只需让 numba 为您初始化 output 数组即可。

我看到的问题是您的行OA = np.empty(dtype=np.float64) 。 在这里你想分配一个数组，但你没有提供形状。 您可以使用方便的 function OA = np.empty_like(Y_LAT) ，它将分配一个大小和数据类型为Y_LAT的空数组。 然而，看看你的下一行（这是一个赋值并丢弃之前分配的数组），你不需要手动分配数组，因为 numba 会为你完成这项工作。 您的（未提供）function nan2zero的参数np.arctan2(Y_LAT, Y_LON) - np.arctan2( X_LAT, X_LON )将已经分配一个数组并用给定计算的 output 填充该数组。 请注意，在您的助手 function coord_unit_vec中，您也不需要手动分配和初始化lat_vec和lon_vec 。 这是由 numba 为您完成的。

因此，删除OA = np.empty(dtype=np.float64)行，也许为了更好的可读性，将下一行拆分为

[...]
# the mathematical expression which defines the value and does the allocation
OA = np.arctan2(Y_LAT, Y_LON) - np.arctan2( X_LAT, X_LON )
# some polishing for the numerics, e.g. removing NaNs.
[...]