Experiments with numba, decorators and closures - poliastro/poliastro GitHub Wiki

In [1]: from poliastro.twobody.decorators import state_from_vector

In [2]: def fun(t, ss):
   ...:     return ss.r
   ...: 
   ...: 

In [3]: f = state_from_vector(fun)

In [4]: f(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
Out[4]: <Quantity [1., 0., 0.] km>

In [6]: import numba

In [7]: numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-5c08b2d75712> in <module>()
----> 1 numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)

TypeError: too many arguments: expected 2, got 3

In [8]: state_from_vector??

In [9]: numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-3714a6a93b81> in <module>()
----> 1 numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)

TypeError: too many arguments: expected 2, got 3

In [10]: numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-10-535ee916f529> in <module>()
----> 1 numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0])

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    358                     e.patch_message(''.join(e.args) + help_msg)
    359             # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 360             raise e
    361 
    362     def inspect_llvm(self, signature=None):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    309                 argtypes.append(self.typeof_pyval(a))
    310         try:
--> 311             return self.compile(tuple(argtypes))
    312         except errors.TypingError as e:
    313             # Intercept typing error that may be due to an argument

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in compile(self, sig)
    616 
    617                 self._cache_misses[sig] += 1
--> 618                 cres = self._compiler.compile(args, return_type)
    619                 self.add_overload(cres)
    620                 self._cache.save_overload(sig, cres)

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in compile(self, args, return_type)
     81                                       args=args, return_type=return_type,
     82                                       flags=flags, locals=self.locals,
---> 83                                       pipeline_class=self.pipeline_class)
     84         # Check typing error if object mode is used
     85         if cres.typing_error is not None and not flags.enable_pyobject:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    869     pipeline = pipeline_class(typingctx, targetctx, library,
    870                               args, return_type, flags, locals)
--> 871     return pipeline.compile_extra(func)
    872 
    873 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in compile_extra(self, func)
    363         self.lifted = ()
    364         self.lifted_from = None
--> 365         return self._compile_bytecode()
    366 
    367     def compile_ir(self, func_ir, lifted=(), lifted_from=None):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in _compile_bytecode(self)
    800         """
    801         assert self.func_ir is None
--> 802         return self._compile_core()
    803 
    804     def _compile_ir(self):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in _compile_core(self)
    787         self.define_pipelines(pm)
    788         pm.finalize()
--> 789         res = pm.run(self.status)
    790         if res is not None:
    791             # Early pipeline completion

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    249                     # No more fallback pipelines?
    250                     if is_final_pipeline:
--> 251                         raise patched_exception
    252                     # Go to next fallback pipeline
    253                     else:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    241                 try:
    242                     event(stage_name)
--> 243                     stage()
    244                 except _EarlyPipelineCompletion as e:
    245                     return e.result

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in stage_analyze_bytecode(self)
    377         Analyze bytecode and translating to Numba IR
    378         """
--> 379         func_ir = translate_stage(self.func_id, self.bc)
    380         self._set_and_check_ir(func_ir)
    381 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in translate_stage(func_id, bytecode)
    933 def translate_stage(func_id, bytecode):
    934     interp = interpreter.Interpreter(func_id)
--> 935     return interp.interpret(bytecode)
    936 
    937 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in interpret(self, bytecode)
     96         # Data flow analysis
     97         self.dfa = dataflow.DataFlowAnalysis(self.cfa)
---> 98         self.dfa.run()
     99 
    100         # Temp states during interpretation

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dataflow.py in run(self)
     26     def run(self):
     27         for blk in self.cfa.iterliveblocks():
---> 28             self.infos[blk.offset] = self.run_on_block(blk)
     29 
     30     def run_on_block(self, blk):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dataflow.py in run_on_block(self, blk)
     76         for offset in blk:
     77             inst = self.bytecode[offset]
---> 78             self.dispatch(info, inst)
     79         return info
     80 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dataflow.py in dispatch(self, info, inst)
     86         fname = "op_%s" % inst.opname.replace('+', '_')
     87         fn = getattr(self, fname, self.handle_unknown_opcode)
---> 88         fn(info, inst)
     89 
     90     def handle_unknown_opcode(self, info, inst):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dataflow.py in op_CALL_FUNCTION_EX(self, info, inst)
    352             if inst.arg & 1:
    353                 errmsg = 'CALL_FUNCTION_EX with **kwargs not supported'
--> 354                 raise NotImplementedError(errmsg)
    355             vararg = info.pop()
    356             func = info.pop()

NotImplementedError: Failed at nopython (analyzing bytecode)
CALL_FUNCTION_EX with **kwargs not supported

A possible issue with numba and decorators?

In [28]: def state_from_vector_simple(func):
    ...:     @wraps(func)
    ...:     def wrapper(t, u_, k):
    ...:         r, v = u_[:3], u_[3:]
    ...:         ss = RVState(Body(None, k * u.km3s2, "_Dummy"), r * u.km, v * u.kms)
    ...: 
    ...:         return func(t, ss)
    ...: 
    ...:     return wrapper
    ...: 
    ...: 

In [29]: 

In [29]: f = state_from_vector_simple(fun)

In [30]: f(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
Out[30]: <Quantity [1., 0., 0.] km>

In [31]: numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-31-5c08b2d75712> in <module>()
----> 1 numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)

TypeError: too many arguments: expected 2, got 3

In [32]: numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-32-3714a6a93b81> in <module>()
----> 1 numba.njit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 1.0)

TypeError: too many arguments: expected 2, got 3

In [33]: numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    113         try:
--> 114             return self._con[name]
    115         except KeyError:

KeyError: 'u_'

During handling of the above exception, another exception occurred:

NotDefinedError                           Traceback (most recent call last)
~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get_exact(self, name)
    718         try:
--> 719             return self.localvars.get(name)
    720         except NotDefinedError:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    115         except KeyError:
--> 116             raise NotDefinedError(name)
    117 

NotDefinedError: 'u_' is not defined in None

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    113         try:
--> 114             return self._con[name]
    115         except KeyError:

KeyError: 'u_'

During handling of the above exception, another exception occurred:

NotDefinedError                           Traceback (most recent call last)
<ipython-input-33-281715f4a79a> in <module>()
----> 1 numba.jit(f)(0.0, [1.0, 0.0, 0.0, 0.0, 1.0, 0.0])

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    358                     e.patch_message(''.join(e.args) + help_msg)
    359             # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 360             raise e
    361 
    362     def inspect_llvm(self, signature=None):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    309                 argtypes.append(self.typeof_pyval(a))
    310         try:
--> 311             return self.compile(tuple(argtypes))
    312         except errors.TypingError as e:
    313             # Intercept typing error that may be due to an argument

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in compile(self, sig)
    616 
    617                 self._cache_misses[sig] += 1
--> 618                 cres = self._compiler.compile(args, return_type)
    619                 self.add_overload(cres)
    620                 self._cache.save_overload(sig, cres)

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/dispatcher.py in compile(self, args, return_type)
     81                                       args=args, return_type=return_type,
     82                                       flags=flags, locals=self.locals,
---> 83                                       pipeline_class=self.pipeline_class)
     84         # Check typing error if object mode is used
     85         if cres.typing_error is not None and not flags.enable_pyobject:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    869     pipeline = pipeline_class(typingctx, targetctx, library,
    870                               args, return_type, flags, locals)
--> 871     return pipeline.compile_extra(func)
    872 
    873 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in compile_extra(self, func)
    363         self.lifted = ()
    364         self.lifted_from = None
--> 365         return self._compile_bytecode()
    366 
    367     def compile_ir(self, func_ir, lifted=(), lifted_from=None):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in _compile_bytecode(self)
    800         """
    801         assert self.func_ir is None
--> 802         return self._compile_core()
    803 
    804     def _compile_ir(self):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in _compile_core(self)
    787         self.define_pipelines(pm)
    788         pm.finalize()
--> 789         res = pm.run(self.status)
    790         if res is not None:
    791             # Early pipeline completion

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    249                     # No more fallback pipelines?
    250                     if is_final_pipeline:
--> 251                         raise patched_exception
    252                     # Go to next fallback pipeline
    253                     else:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    241                 try:
    242                     event(stage_name)
--> 243                     stage()
    244                 except _EarlyPipelineCompletion as e:
    245                     return e.result

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in stage_analyze_bytecode(self)
    377         Analyze bytecode and translating to Numba IR
    378         """
--> 379         func_ir = translate_stage(self.func_id, self.bc)
    380         self._set_and_check_ir(func_ir)
    381 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/compiler.py in translate_stage(func_id, bytecode)
    933 def translate_stage(func_id, bytecode):
    934     interp = interpreter.Interpreter(func_id)
--> 935     return interp.interpret(bytecode)
    936 
    937 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in interpret(self, bytecode)
    108         # Interpret loop
    109         for inst, kws in self._iter_inst():
--> 110             self._dispatch(inst, kws)
    111 
    112         return ir.FunctionIR(self.blocks, self.is_generator, self.func_id,

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in _dispatch(self, inst, kws)
    237                 if e.loc is None:
    238                     e.loc = self.loc
--> 239                 raise e
    240 
    241 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in _dispatch(self, inst, kws)
    233         else:
    234             try:
--> 235                 return fn(inst, **kws)
    236             except errors.NotDefinedError as e:
    237                 if e.loc is None:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in op_LOAD_FAST(self, inst, res)
    538     def op_LOAD_FAST(self, inst, res):
    539         srcname = self.code_locals[inst.arg]
--> 540         self.store(value=self.get(srcname), name=res)
    541 
    542     def op_STORE_FAST(self, inst, value):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/interpreter.py in get(self, name)
    271         var = self.assigner.get_assignment_source(name)
    272         if var is None:
--> 273             var = self.current_scope.get(name)
    274         return var
    275 

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    709         if name in self.redefined:
    710             name = "%s.%d" % (name, self.redefined[name])
--> 711         return self.get_exact(name)
    712 
    713     def get_exact(self, name):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get_exact(self, name)
    720         except NotDefinedError:
    721             if self.has_parent:
--> 722                 return self.parent.get(name)
    723             else:
    724                 raise

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    709         if name in self.redefined:
    710             name = "%s.%d" % (name, self.redefined[name])
--> 711         return self.get_exact(name)
    712 
    713     def get_exact(self, name):

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get_exact(self, name)
    717         """
    718         try:
--> 719             return self.localvars.get(name)
    720         except NotDefinedError:
    721             if self.has_parent:

~/.miniconda36/envs/poliastro36/lib/python3.6/site-packages/numba/ir.py in get(self, name)
    114             return self._con[name]
    115         except KeyError:
--> 116             raise NotDefinedError(name)
    117 
    118     def __contains__(self, name):

NotDefinedError: 'u_' is not defined in <ipython-input-28-1fe7129d9a2f> (4)
⚠️ **GitHub.com Fallback** ⚠️