Python - Numba 产生 For-Loop 的语法错误

Question

I have the following code in which I create two 1D NumPy arrays (f0src and f1src) and then create a two for loop for calculate some function with every combination of f0src and f1src arrays. I pre-allocate the output to fasten the process (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)

This works fine and gives me the output I desire but it is slow and with relatively a small sample array it takes around 1:35 second

CPU times: user 1min 34s, sys: 564 ms, total: 1min 34s Wall time: 1min 35s

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

Based on what I read, Numba can fasten the process quite much, especially if you are using NumPy arrays, you pre-allocate output and have for loops which is exactly my case. So I tried the same thing by just adding njit before the function

@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)

However, I am getting the following error that I don't fully understand.

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3444, in run_code exec(code_obj, self.user_global_ns, self.user_ns)

File "/var/folders/y9/nvl5y5_15v7cx8wb6nv3lzxc0000gn/T/ipykernel_17771/4021014180.py", line 1, in 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')

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 2406, in run_cell_magic result = fn(*args, **kwargs)

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/decorator.py", line 232, in fun return caller(func, *(extras + args), **kw)

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/magic.py", line 187, in call = lambda f, *a, **k: f(*a, **k)

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/magics/execution.py", line 1280, in time expr_ast = self.shell.compile.ast_parse(expr)

File "/Users/sara/opt/miniconda3/lib/python3.9/site-packages/IPython/core/compilerop.py", line 101, in ast_parse return compile(source, filename, symbol, self.flags | PyCF_ONLY_AST, 1)

File "", line 2 for idxf0, f0 in enumerate(f0src): ^ SyntaxError: invalid syntax

I appreciate it if you could point me in the right direction.

Answer 1

@njit is a function decorator. This means that it should wrap some function, and, thus, you should provide it a function, not loop. Use the following:

@njit    
def foo(mytime, f0src, f1src, z3_2d):
   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, z3_2d)

This code creates some function foo with operations you need, and uses @njit decorator to wrap it. The last argument is mutable, so it will store the result as expected. Notice that you should also use @njit decorator for other functions in foo (my_phase, z_n).