python - 由于数据类型，编译 njit nopython 版本的函数失败

Question

我正在 njit 中编写一个函数来加速非常慢的水库操作优化代码。该函数根据水库水位和闸门可用性返回泄漏释放的最大值。我传入一个参数大小，指定要计算的流数(在某些调用中它是一个，在某些调用中它是多个)。我还传递了一个 numpy.zeros 数组，然后我可以用函数输出填充该数组。该函数的简化版本编写如下:

import numpy as np
from numba import njit

@njit(cache=True)
def fncMaxFlow(elev, flag, size, MaxQ):
    if (flag == 1): # SPOG2 running
        if size==0:
            if (elev>367.28):
                return 861.1 
            else: return 0
        else:
            for i in range(size):
                if((elev[i]>367.28) & (elev[i]<385)):
                    MaxQ[i]=861.1
            return MaxQ
    else:
        if size==0: return 0
        else: return MaxQ

fncMaxFlow(np.random.randint(368, 380, 3), 1, 3, np.zeros(3))

我收到的错误:

Can't unify return type from the following types: array(float64, 1d, C), float64, int32

这是什么原因呢？是否有任何解决方法或我缺少的某些步骤，以便我可以使用 numba 来加快速度？该函数和其他类似函数被调用数百万次，因此它们是计算效率的主要因素。任何建议都会有所帮助 - 我对 python 还很陌生。

Answer 1

numba 函数中的变量必须具有一致的类型，包括返回变量。在代码中，您可以返回 MaxQ(数组)、861.1( float )或 0(整数)。

您需要重构此代码，以便无论代码路径如何，它始终返回一致的类型。

另请注意，在将 numpy 数组与标量 (elev > 367.28) 进行比较的几个地方，您返回的是 bool 值数组，这将导致您问题。因此，您的示例函数不会作为纯 python 函数运行(删除 numba 装饰器)。