Error with mnnCorrect, maybe caused by numba?

Hej all,

I get a weird error when using mnnCorrect on my data. I tried on a simple example, and it fails in exactly the same way. I think it might actually be a problem with numba. Which is the numba version that works with scanpy? (mine is version 0.43.1)

Commands:

pbmc = sc.datasets.pbmc68k_reduced()
sc.external.pp.mnn_correct(pbmc, batch_key="phase")

Here is the error message. As you see it involves a lot of numba warnings:

/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:14: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "l2_norm" failed type inference due to: Invalid use of Function(<function norm at 0x2b3fffedb1e0>) with argument(s) of type(s): (axis=Literal[int](1), x=array(float32, 2d, A))
 * parameterized
In definition 0:
    TypeError: norm_impl() got an unexpected keyword argument 'x'
    raised from /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/typing/templates.py:475
In definition 1:
    TypeError: norm_impl() got an unexpected keyword argument 'x'
    raised from /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/typing/templates.py:475
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<function norm at 0x2b3fffedb1e0>)
[2] During: typing of call at /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py (16)


File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 16:
def l2_norm(in_matrix):
    return np.linalg.norm(x=in_matrix, axis=1)
    ^

  @jit(float32[:](float32[:, :]), nogil=True)
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py:725: NumbaWarning: Function "l2_norm" was compiled in object mode without forceobj=True.

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 15:
@jit(float32[:](float32[:, :]), nogil=True)
def l2_norm(in_matrix):
^

  self.func_ir.loc))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py:734: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 15:
@jit(float32[:](float32[:, :]), nogil=True)
def l2_norm(in_matrix):
^

  warnings.warn(errors.NumbaDeprecationWarning(msg, self.func_ir.loc))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:14: NumbaWarning: Code running in object mode won't allow parallel execution despite nogil=True.
  @jit(float32[:](float32[:, :]), nogil=True)
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:29: NumbaPerformanceWarning: np.dot() is faster on contiguous arrays, called on (array(float32, 1d, A), array(float32, 1d, A))
  dist[i, j] = np.dot(m[i], n[j])
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:88: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "find_mutual_nn" failed type inference due to: Untyped global name 'cKDTree': cannot determine Numba type of <class 'type'>

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 90:
def find_mutual_nn(data1, data2, k1, k2, n_jobs):
    k_index_1 = cKDTree(data1).query(x=data2, k=k1, n_jobs=n_jobs)[1]
    ^

  @jit((float32[:, :], float32[:, :], int8, int8, int8))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:88: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "find_mutual_nn" failed type inference due to: Untyped global name 'cKDTree': cannot determine Numba type of <class 'type'>

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 90:
def find_mutual_nn(data1, data2, k1, k2, n_jobs):
    k_index_1 = cKDTree(data1).query(x=data2, k=k1, n_jobs=n_jobs)[1]
    ^

  @jit((float32[:, :], float32[:, :], int8, int8, int8))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py:725: NumbaWarning: Function "find_mutual_nn" was compiled in object mode without forceobj=True, but has lifted loops.

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 89:
@jit((float32[:, :], float32[:, :], int8, int8, int8))
def find_mutual_nn(data1, data2, k1, k2, n_jobs):
^

  self.func_ir.loc))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py:734: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 89:
@jit((float32[:, :], float32[:, :], int8, int8, int8))
def find_mutual_nn(data1, data2, k1, k2, n_jobs):
^

  warnings.warn(errors.NumbaDeprecationWarning(msg, self.func_ir.loc))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:102: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "compute_correction" failed type inference due to: Invalid use of Function(<function unique at 0x2b3ffff680d0>) with argument(s) of type(s): (array(int32, 1d, A), return_counts=bool)
 * parameterized
In definition 0:
    TypeError: np_unique() got an unexpected keyword argument 'return_counts'
    raised from /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/typing/templates.py:475
In definition 1:
    TypeError: np_unique() got an unexpected keyword argument 'return_counts'
    raised from /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/typing/templates.py:475
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<function unique at 0x2b3ffff680d0>)
[2] During: typing of call at /home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py (105)


File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 105:
def compute_correction(data1, data2, mnn1, mnn2, data2_or_raw2, sigma):
    <source elided>
    vect = data1[mnn1] - data2[mnn2]
    mnn_index, mnn_count = np.unique(mnn2, return_counts=True)
    ^

  @jit(float32[:, :](float32[:, :], float32[:, :], int32[:], int32[:], float32[:, :], float32))
/home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py:102: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "compute_correction" failed type inference due to: cannot determine Numba type of <class 'numba.dispatcher.LiftedLoop'>

File "../../../home/samuele/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py", line 107:
def compute_correction(data1, data2, mnn1, mnn2, data2_or_raw2, sigma):
    <source elided>
    vect_reduced = np.zeros((data2.shape[0], vect.shape[1]), dtype=np.float32)
    for index, ve in zip(mnn2, vect):
    ^

  @jit(float32[:, :](float32[:, :], float32[:, :], int32[:], int32[:], float32[:, :], float32))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-7-341804fa3415> in <module>
      1 pbmc = sc.datasets.pbmc68k_reduced()
----> 2 sc.external.pp.mnn_correct(pbmc, batch_key="phase")

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/scanpy/preprocessing/_mnn_correct.py in mnn_correct(var_index, var_subset, batch_key, index_unique, batch_categories, k, sigma, cos_norm_in, cos_norm_out, svd_dim, var_adj, compute_angle, mnn_order, svd_mode, do_concatenate, save_raw, n_jobs, *datas, **kwargs)
     91     """
     92     try:
---> 93         from mnnpy import mnn_correct as mnn_cor
     94         n_jobs = settings.n_jobs if n_jobs is None else n_jobs
     95         datas, mnn_list, angle_list = mnn_cor(

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/__init__.py in <module>
----> 1 from .mnn import mnn_correct

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/mnn.py in <module>
      5 from anndata import AnnData
      6 from pandas import DataFrame
----> 7 from .utils import transform_input_data, find_mutual_nn, compute_correction
      8 from .utils import svd_internal, find_shared_subspace, get_bio_span, subtract_bio
      9 from .utils import adjust_shift_variance

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/mnnpy/utils.py in <module>
    100 
    101 
--> 102 @jit(float32[:, :](float32[:, :], float32[:, :], int32[:], int32[:], float32[:, :], float32))
    103 def compute_correction(data1, data2, mnn1, mnn2, data2_or_raw2, sigma):
    104     vect = data1[mnn1] - data2[mnn2]

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/decorators.py in wrapper(func)
    184             with typeinfer.register_dispatcher(disp):
    185                 for sig in sigs:
--> 186                     disp.compile(sig)
    187                 disp.disable_compile()
    188         return disp

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler_lock.py in _acquire_compile_lock(*args, **kwargs)
     30         def _acquire_compile_lock(*args, **kwargs):
     31             with self:
---> 32                 return func(*args, **kwargs)
     33         return _acquire_compile_lock
     34 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/dispatcher.py in compile(self, sig)
    657 
    658             self._cache_misses[sig] += 1
--> 659             cres = self._compiler.compile(args, return_type)
    660             self.add_overload(cres)
    661             self._cache.save_overload(sig, cres)

~/miniconda3/envs/scRNA/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:

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    953     pipeline = pipeline_class(typingctx, targetctx, library,
    954                               args, return_type, flags, locals)
--> 955     return pipeline.compile_extra(func)
    956 
    957 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in compile_extra(self, func)
    375         self.lifted = ()
    376         self.lifted_from = None
--> 377         return self._compile_bytecode()
    378 
    379     def compile_ir(self, func_ir, lifted=(), lifted_from=None):

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in _compile_bytecode(self)
    884         """
    885         assert self.func_ir is None
--> 886         return self._compile_core()
    887 
    888     def _compile_ir(self):

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in _compile_core(self)
    871         self.define_pipelines(pm)
    872         pm.finalize()
--> 873         res = pm.run(self.status)
    874         if res is not None:
    875             # Early pipeline completion

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler_lock.py in _acquire_compile_lock(*args, **kwargs)
     30         def _acquire_compile_lock(*args, **kwargs):
     31             with self:
---> 32                 return func(*args, **kwargs)
     33         return _acquire_compile_lock
     34 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    252                     # No more fallback pipelines?
    253                     if is_final_pipeline:
--> 254                         raise patched_exception
    255                     # Go to next fallback pipeline
    256                     else:

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    243                 try:
    244                     event("-- %s" % stage_name)
--> 245                     stage()
    246                 except _EarlyPipelineCompletion as e:
    247                     return e.result

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in stage_objectmode_frontend(self)
    461         if self.flags.enable_looplift:
    462             assert not self.lifted
--> 463             cres = self.frontend_looplift()
    464             if cres is not None:
    465                 raise _EarlyPipelineCompletion(cres)

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in frontend_looplift(self)
    433                               self.args, self.return_type,
    434                               outer_flags, self.locals,
--> 435                               lifted=tuple(loops), lifted_from=None)
    436             return cres
    437 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in compile_ir(typingctx, targetctx, func_ir, args, return_type, flags, locals, lifted, lifted_from, is_lifted_loop, library, pipeline_class)
   1015                                   args, return_type, flags, locals)
   1016         return pipeline.compile_ir(func_ir=func_ir, lifted=lifted,
-> 1017                                 lifted_from=lifted_from)
   1018 
   1019 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in compile_ir(self, func_ir, lifted, lifted_from)
    383 
    384         self._set_and_check_ir(func_ir)
--> 385         return self._compile_ir()
    386 
    387     def stage_analyze_bytecode(self):

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in _compile_ir(self)
    891         """
    892         assert self.func_ir is not None
--> 893         return self._compile_core()
    894 
    895 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in _compile_core(self)
    871         self.define_pipelines(pm)
    872         pm.finalize()
--> 873         res = pm.run(self.status)
    874         if res is not None:
    875             # Early pipeline completion

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler_lock.py in _acquire_compile_lock(*args, **kwargs)
     30         def _acquire_compile_lock(*args, **kwargs):
     31             with self:
---> 32                 return func(*args, **kwargs)
     33         return _acquire_compile_lock
     34 

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    252                     # No more fallback pipelines?
    253                     if is_final_pipeline:
--> 254                         raise patched_exception
    255                     # Go to next fallback pipeline
    256                     else:

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in run(self, status)
    243                 try:
    244                     event("-- %s" % stage_name)
--> 245                     stage()
    246                 except _EarlyPipelineCompletion as e:
    247                     return e.result

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in stage_objectmode_backend(self)
    709         """
    710         lowerfn = self.backend_object_mode
--> 711         self._backend(lowerfn, objectmode=True)
    712 
    713         # Warn, deprecated behaviour, code compiled in objmode without

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in _backend(self, lowerfn, objectmode)
    685             self.library.enable_object_caching()
    686 
--> 687         lowered = lowerfn()
    688         signature = typing.signature(self.return_type, *self.args)
    689         self.cr = compile_result(

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in backend_object_mode(self)
    657                                      self.library,
    658                                      self.func_ir,
--> 659                                      self.flags)
    660 
    661     def backend_nopython_mode(self):

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/compiler.py in py_lowering_stage(targetctx, library, interp, flags)
   1146     with targetctx.push_code_library(library):
   1147         lower = pylowering.PyLower(targetctx, library, fndesc, interp)
-> 1148         lower.lower()
   1149         if not flags.no_cpython_wrapper:
   1150             lower.create_cpython_wrapper()

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/lowering.py in lower(self)
    200 
    201         # Materialize LLVM Module
--> 202         self.library.add_ir_module(self.module)
    203 
    204     def extract_function_arguments(self):

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/numba/targets/codegen.py in add_ir_module(self, ir_module)
    183         assert isinstance(ir_module, llvmir.Module)
    184         ir = cgutils.normalize_ir_text(str(ir_module))
--> 185         ll_module = ll.parse_assembly(ir)
    186         ll_module.name = ir_module.name
    187         ll_module.verify()

~/miniconda3/envs/scRNA/lib/python3.6/site-packages/llvmlite/binding/module.py in parse_assembly(llvmir, context)
     24         if errmsg:
     25             mod.close()
---> 26             raise RuntimeError("LLVM IR parsing error\n{0}".format(errmsg))
     27     return mod
     28 

RuntimeError: Failed in object mode pipeline (step: object mode frontend)
Failed in object mode pipeline (step: object mode backend)
LLVM IR parsing error
<string>:2792:50: error: expected type
  %".2433" = call i8* @"PyLong_FromSsize_t"(i64 (slice(None, None, None), None))
                                                 ^

Cheers,
Samuele

Hi Samuele,

This might be a shot in the dark, but I was under the impression that sc.external.pp.mnn_correct() requires separate datasets as input. If that is still the case, then you would have to first split the pbmc datasets by phase before putting them into sc.external.pp.mnn_correct().

So do something like:

datasets = [pbmc[pbmc.obs['phase'] == ph].copy() for ph in pbmc.obs.phase.cat.categories]
sc.external.pp.mnn_correct(*datasets)

Both way did not work. I had indeed to downgrande numba below version 0.42!

The function runs for me but I’m gettin the same fall back compilation warning (but no runtime error). scanpy==1.4.6, mnnpy==0.1.9.5 and numba 0.48.0. Any idea which numba version matches these versions of mnnpy and scanpy? Or is there something else I’m missing… Concerns me that the batch correction isn’t working effectively. Hard to tell with the plots of the uncorrected vs corrected datasets (both subsetted with HVGs and using same PCs) afterward.