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)