Python - Numba 为 For-Loop 产生语法错误答案

【问题标题】：Python - Numba produces syntax error for For-LoopPython - Numba 为 For-Loop 产生语法错误
【发布时间】：2022-08-20 17:08:22
【问题描述】：

我有以下代码，我在其中创建了两个 1D NumPy 数组（f0src 和 f1src），然后创建了一个两个 for 循环，以使用 f0src 和 f1src 数组的每个组合计算某个函数。我预先分配输出以加快进程（z3_2d）。

n   = 50
f0  = 5.073174792651648
f1  = -1.50272e-13
df0  = 1e-7 
df1 = 8e-15

f0src = np.arange(f0 - n * df0, f0 + n * df0, df0)
f1src = np.arange(f1 - n * df1, f1 + n * df1, df1)

f0shape=f0src.shape[0]
f1shape=f1src.shape[0]
z3_2d = np.zeros([f0shape,f1shape])

%%time

for idxf0, f0 in enumerate(f0src):

    for idxf1, f1 in enumerate(f1src):

        phase=my_phase(mytime,f0,f1) #mytime is another 1D array around 100k
        z3=z_n(phase, n=3, norm=1)
        z3_2d[idxf0, idxf1]=np.copy(z3)

这工作正常，给了我想要的输出，但它很慢，并且样本数组相对较小，大约需要 1:35 秒

CPU时间：用户1min 34s，sys：564 ms，总计：1min 34s Wall time：1min 35秒

z3_2d.shape ---> (100, 100)

根据我阅读的内容，努巴可以大大加快这个过程，特别是如果你使用 NumPy 数组，你会预先分配输出并有 for 循环，这正是我的情况。所以我尝试了同样的事情，只需添加njit在函数之前
@njit
for idxf0, f0 in enumerate(f0src):

    for idxf1, f1 in enumerate(f1src):

        phase=my_phase(mytime,f0,f1)
        z3=z_n(phase, n=3, norm=1)
        z3_2d[idxf0, idxf1]=np.copy(z3)
但是，我收到以下我不完全理解的错误。

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/interactiveshell.py\", 第 3444 行，在 run_code 中执行（code_obj，self.user_global_ns，self.user_ns）

文件 \"/var/folders/y9/nvl5y5_15v7cx8wb6nv3lzxc0000gn/T/ipykernel_17771/4021014180.py\", 第 1 行，在 get_ipython().run_cell_magic(\'time\', \'\', \'\\n@njit\\nfor idxf0, f0 in enumerate(f0src):\\n\\n for idxf1, f1 in enumerate(f1src ):\\n\\n
phase=pulse_phase(timemerged,f0,f1)\\n z3=z_n(phase, n=3, norm=1)\\n z3_2d[idxf0, idxf1]=np.copy(z3)\\n\')

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/interactiveshell.py\", 第 2406 行，在 run_cell_magic 结果 = fn(*args, **kwargs)

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/decorator.py\", 第 232 行，有趣返回调用者(func, *(extras + args), **kw)

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/magic.py\", 第 187 行，在调用 = lambda f, *a, **k: f(*a, **k)

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/magics/execution.py\", 第 1280 行，及时 expr_ast = self.shell.compile.ast_parse(expr)

文件 \"/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/compilerop.py\", 第 101 行，在 ast_parse 中返回编译（源，文件名，符号，self.flags | PyCF_ONLY_AST，1）

文件 \"\"，第 2 行对于枚举（f0src）中的 idxf0，f0： ^ SyntaxError: 无效的语法

如果您能指出正确的方向，我将不胜感激。

我认为您需要将两个 for 循环放在一个实际的函数中，然后在该函数上使用 @njit 包装器。
my_phase 和 z_n 是 python 函数，所以我不知道你期望 numba 在这里做什么。

标签： python arrays python-3.x numpy numba

【解决方案1】：

@njit 是一个函数装饰器。这意味着它应该包装一些函数，因此，您应该为它提供一个函数，而不是循环。使用以下内容：

@njit    
def foo(mytime, f0src, f1src):
   for idxf0, f0 in enumerate(f0src):

      for idxf1, f1 in enumerate(f1src):

         phase=my_phase(mytime,f0,f1)
         z3=z_n(phase, n=3, norm=1)
         z3_2d[idxf0, idxf1]=np.copy(z3)
foo(mytime, f0src, f1src)

请注意，您还应该将@njit 装饰器用于此功能中的其他功能（my_phase，z_n）

【讨论】：