Compute Transition Matrix Error CellRank Meets RNA Velocity #1211

anthonyinsalaco1 · 2024-07-15T16:46:03Z

When I try to run the code below on the tutorial, I get this error.

```Computing transition matrix using `'deterministic'` model
0%
0/2531 [00:00<?, ?cell/s]

KeyError Traceback (most recent call last)
File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/caching.py:469, in IndexDataCacheFile.save(self, key, data)
467 try:
468 # If key already exists, we will overwrite the file
--> 469 data_name = overloads[key]
470 except KeyError:
471 # Find an available name for the data file

KeyError: ((Function(<function norm at 0x7f7d371c2af0>), Literalint, Array(float64, 2, 'A', False, aligned=True)), ('x86_64-apple-darwin23.4.0', 'skylake', '+64bit,+adx,+aes,-amx-bf16,-amx-int8,-amx-tile,+avx,+avx2,-avx512bf16,-avx512bitalg,-avx512bw,-avx512cd,-avx512dq,-avx512er,-avx512f,-avx512fp16,-avx512ifma,-avx512pf,-avx512vbmi,-avx512vbmi2,-avx512vl,-avx512vnni,-avx512vp2intersect,-avx512vpopcntdq,-avxvnni,+bmi,+bmi2,-cldemote,+clflushopt,-clwb,-clzero,+cmov,+crc32,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,-gfni,-hreset,+invpcid,-kl,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mwaitx,+pclmul,-pconfig,-pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,-rdpid,+rdrnd,+rdseed,-rtm,+sahf,-serialize,+sgx,-sha,-shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,-sse4a,+ssse3,-tbm,-tsxldtrk,-uintr,-vaes,-vpclmulqdq,-waitpkg,-wbnoinvd,-widekl,-xop,+xsave,+xsavec,+xsaveopt,+xsaves'), ('04d17047dec8a5ada0d1f7224d7e8baf78f8fb78c8913cbdf8a204a5c3f1492f', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'))

During handling of the above exception, another exception occurred:

TypeError Traceback (most recent call last)
Input In [23], in <cell line: 1>()
----> 1 vk.compute_transition_matrix()

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/_velocity_kernel.py:160, in VelocityKernel.compute_transition_matrix(self, model, backward_mode, similarity, softmax_scale, n_samples, seed, **kwargs)
157 return self
159 if softmax_scale is None:
--> 160 softmax_scale = self._estimate_softmax_scale(backward_mode=backward_mode, similarity=similarity)
161 logg.info(f"Using softmax_scale={softmax_scale:.4f}")
162 params["softmax_scale"] = softmax_scale

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/_velocity_kernel.py:254, in VelocityKernel._estimate_softmax_scale(self, n_jobs, backend, **kwargs)
247 def _estimate_softmax_scale(
248 self,
249 n_jobs: Optional[int] = None,
250 backend: Backend_t = DEFAULT_BACKEND,
251 **kwargs,
252 ) -> float:
253 model = self._create_model(VelocityModel.DETERMINISTIC, softmax_scale=1.0, **kwargs)
--> 254 _, logits = model(n_jobs, backend)
255 return 1.0 / np.median(np.abs(logits.data))

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:64, in ModelABC.call(self, n_jobs, backend, show_progress_bar, **kwargs)
56 def call(
57 self,
58 n_jobs: Optional[int] = None,
(...)
61 **kwargs: Any,
62 ) -> Tuple[np.ndarray, np.ndarray]:
63 ixs = self._ixs
---> 64 return parallelize(
65 self._compute_helper,
66 ixs,
67 n_jobs=n_jobs,
68 backend=backend,
69 show_progress_bar=show_progress_bar,
70 as_array=False,
71 extractor=lambda data: self._reconstruct_output(np.concatenate(data, axis=-1), ixs),
72 unit=self._unit,
73 )(**kwargs)

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/_utils/_parallelize.py:96, in parallelize..wrapper(*args, **kwargs)
93 else:
94 pbar, queue, thread = None, None, None
---> 96 res = jl.Parallel(n_jobs=n_jobs, backend=backend)(
97 jl.delayed(callback)(
98 *((i, cs) if use_ixs else (cs,)),
99 *args,
100 **kwargs,
101 queue=queue,
102 )
103 for i, cs in enumerate(collections)
104 )
106 res = np.array(res) if as_array else res
107 if thread is not None:

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:1043, in Parallel.call(self, iterable)
1034 try:
1035 # Only set self._iterating to True if at least a batch
1036 # was dispatched. In particular this covers the edge
(...)
1040 # was very quick and its callback already dispatched all the
1041 # remaining jobs.
1042 self._iterating = False
-> 1043 if self.dispatch_one_batch(iterator):
1044 self._iterating = self._original_iterator is not None
1046 while self.dispatch_one_batch(iterator):

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:861, in Parallel.dispatch_one_batch(self, iterator)
859 return False
860 else:
--> 861 self._dispatch(tasks)
862 return True

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:779, in Parallel._dispatch(self, batch)
777 with self._lock:
778 job_idx = len(self._jobs)
--> 779 job = self._backend.apply_async(batch, callback=cb)
780 # A job can complete so quickly than its callback is
781 # called before we get here, causing self._jobs to
782 # grow. To ensure correct results ordering, .insert is
783 # used (rather than .append) in the following line
784 self._jobs.insert(job_idx, job)

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/_parallel_backends.py:208, in SequentialBackend.apply_async(self, func, callback)
206 def apply_async(self, func, callback=None):
207 """Schedule a func to be run"""
--> 208 result = ImmediateResult(func)
209 if callback:
210 callback(result)

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/_parallel_backends.py:572, in ImmediateResult.init(self, batch)
569 def init(self, batch):
570 # Don't delay the application, to avoid keeping the input
571 # arguments in memory
--> 572 self.results = batch()

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:262, in BatchedCalls.call(self)
258 def call(self):
259 # Set the default nested backend to self._backend but do not set the
260 # change the default number of processes to -1
261 with parallel_backend(self._backend, n_jobs=self._n_jobs):
--> 262 return [func(*args, **kwargs)
263 for func, args, kwargs in self.items]

File ~/opt/anaconda3/lib/python3.9/site-packages/joblib/parallel.py:262, in (.0)
258 def call(self):
259 # Set the default nested backend to self._backend but do not set the
260 # change the default number of processes to -1
261 with parallel_backend(self._backend, n_jobs=self._n_jobs):
--> 262 return [func(*args, **kwargs)
263 for func, args, kwargs in self.items]

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:85, in ModelABC._compute_helper(self, ixs, queue, **kwargs)
82 neigh_ixs = indices[start:end]
83 n_neigh = len(neigh_ixs)
---> 85 ps, ls = self._compute(ix, neigh_ixs, **kwargs)
86 if np.shape(ps) != (n_neigh,):
87 raise ValueError(f"Expected row of shape {(2, n_neigh)}, found {np.shape(ps)}.")

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_velocity_model.py:181, in Deterministic._compute(self, ix, neigh_ixs)
178 if self._backward_mode == BackwardMode.NEGATE:
179 v *= -1.0
--> 181 return self._similarity(v[None, :], W, self._softmax_scale)

File ~/opt/anaconda3/lib/python3.9/site-packages/cellrank/kernels/utils/_similarity.py:238, in SimilarityHessian.call(self, v, D, softmax_scale)
226 @d.dedent
227 def call(self, v: np.ndarray, D: np.ndarray, softmax_scale: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
228 """%(sim_scheme.full_desc)s
229
230 Parameters
(...)
236 %(sim_scheme.returns)s
237 """ # noqa: D400
--> 238 return _predict_transition_probabilities_numpy(v, D, softmax_scale, self._center_mean, self._scale_by_norm)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:442, in _DispatcherBase._compile_for_args(self, *args, **kws)
440 e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
441 # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 442 raise e
443 finally:
444 self._types_active_call = []

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:375, in _DispatcherBase._compile_for_args(self, *args, **kws)
373 return_val = None
374 try:
--> 375 return_val = self.compile(tuple(argtypes))
376 except errors.ForceLiteralArg as e:
377 # Received request for compiler re-entry with the list of arguments
378 # indicated by e.requested_args.
379 # First, check if any of these args are already Literal-ized
380 already_lit_pos = [i for i in e.requested_args
381 if isinstance(args[i], types.Literal)]

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:905, in Dispatcher.compile(self, sig)
903 with ev.trigger_event("numba:compile", data=ev_details):
904 try:
--> 905 cres = self._compiler.compile(args, return_type)
906 except errors.ForceLiteralArg as e:
907 def folded(args, kws):

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:80, in _FunctionCompiler.compile(self, args, return_type)
79 def compile(self, args, return_type):
---> 80 status, retval = self._compile_cached(args, return_type)
81 if status:
82 return retval

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:94, in _FunctionCompiler._compile_cached(self, args, return_type)
91 pass
93 try:
---> 94 retval = self._compile_core(args, return_type)
95 except errors.TypingError as e:
96 self._failed_cache[key] = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:107, in _FunctionCompiler._compile_core(self, args, return_type)
104 flags = self._customize_flags(flags)
106 impl = self._get_implementation(args, {})
--> 107 cres = compiler.compile_extra(self.targetdescr.typing_context,
108 self.targetdescr.target_context,
109 impl,
110 args=args, return_type=return_type,
111 flags=flags, locals=self.locals,
112 pipeline_class=self.pipeline_class)
113 # Check typing error if object mode is used
114 if cres.typing_error is not None and not flags.enable_pyobject:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:744, in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
720 """Compiler entry point
721
722 Parameter
(...)
740 compiler pipeline
741 """
742 pipeline = pipeline_class(typingctx, targetctx, library,
743 args, return_type, flags, locals)
--> 744 return pipeline.compile_extra(func)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:438, in CompilerBase.compile_extra(self, func)
436 self.state.lifted = ()
437 self.state.lifted_from = None
--> 438 return self._compile_bytecode()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:506, in CompilerBase._compile_bytecode(self)
502 """
503 Populate and run pipeline for bytecode input
504 """
505 assert self.state.func_ir is None
--> 506 return self._compile_core()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:481, in CompilerBase._compile_core(self)
478 except Exception as e:
479 if (utils.use_new_style_errors() and not
480 isinstance(e, errors.NumbaError)):
--> 481 raise e
483 self.state.status.fail_reason = e
484 if is_final_pipeline:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler.py:472, in CompilerBase._compile_core(self)
470 res = None
471 try:
--> 472 pm.run(self.state)
473 if self.state.cr is not None:
474 break

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:364, in PassManager.run(self, state)
361 except Exception as e:
362 if (utils.use_new_style_errors() and not
363 isinstance(e, errors.NumbaError)):
--> 364 raise e
365 msg = "Failed in %s mode pipeline (step: %s)" %
366 (self.pipeline_name, pass_desc)
367 patched_exception = self._patch_error(msg, e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:356, in PassManager.run(self, state)
354 pass_inst = _pass_registry.get(pss).pass_inst
355 if isinstance(pass_inst, CompilerPass):
--> 356 self._runPass(idx, pass_inst, state)
357 else:
358 raise BaseException("Legacy pass in use")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.call.._acquire_compile_lock(*args, **kwargs)
32 @functools.wraps(func)
33 def _acquire_compile_lock(*args, **kwargs):
34 with self:
---> 35 return func(*args, **kwargs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:311, in PassManager._runPass(self, index, pss, internal_state)
309 mutated |= check(pss.run_initialization, internal_state)
310 with SimpleTimer() as pass_time:
--> 311 mutated |= check(pss.run_pass, internal_state)
312 with SimpleTimer() as finalize_time:
313 mutated |= check(pss.run_finalizer, internal_state)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/compiler_machinery.py:273, in PassManager._runPass..check(func, compiler_state)
272 def check(func, compiler_state):
--> 273 mangled = func(compiler_state)
274 if mangled not in (True, False):
275 msg = ("CompilerPass implementations should return True/False. "
276 "CompilerPass with name '%s' did not.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:112, in BaseTypeInference.run_pass(self, state)
106 """
107 Type inference and legalization
108 """
109 with fallback_context(state, 'Function "%s" failed type inference'
110 % (state.func_id.func_name,)):
111 # Type inference
--> 112 typemap, return_type, calltypes, errs = type_inference_stage(
113 state.typingctx,
114 state.targetctx,
115 state.func_ir,
116 state.args,
117 state.return_type,
118 state.locals,
119 raise_errors=self._raise_errors)
120 state.typemap = typemap
121 # save errors in case of partial typing

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typed_passes.py:93, in type_inference_stage(typingctx, targetctx, interp, args, return_type, locals, raise_errors)
91 infer.build_constraint()
92 # return errors in case of partial typing
---> 93 errs = infer.propagate(raise_errors=raise_errors)
94 typemap, restype, calltypes = infer.unify(raise_errors=raise_errors)
96 return _TypingResults(typemap, restype, calltypes, errs)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1083, in TypeInferer.propagate(self, raise_errors)
1080 oldtoken = newtoken
1081 # Errors can appear when the type set is incomplete; only
1082 # raise them when there is no progress anymore.
-> 1083 errors = self.constraints.propagate(self)
1084 newtoken = self.get_state_token()
1085 self.debug.propagate_finished()

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:182, in ConstraintNetwork.propagate(self, typeinfer)
180 errors.append(utils.chain_exception(new_exc, e))
181 elif utils.use_new_style_errors():
--> 182 raise e
183 else:
184 msg = ("Unknown CAPTURED_ERRORS style: "
185 f"'{config.CAPTURED_ERRORS}'.")

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:160, in ConstraintNetwork.propagate(self, typeinfer)
157 with typeinfer.warnings.catch_warnings(filename=loc.filename,
158 lineno=loc.line):
159 try:
--> 160 constraint(typeinfer)
161 except ForceLiteralArg as e:
162 errors.append(e)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:583, in CallConstraint.call(self, typeinfer)
581 fnty = typevars[self.func].getone()
582 with new_error_context("resolving callee type: {0}", fnty):
--> 583 self.resolve(typeinfer, typevars, fnty)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:606, in CallConstraint.resolve(self, typeinfer, typevars, fnty)
604 fnty = fnty.instance_type
605 try:
--> 606 sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
607 except ForceLiteralArg as e:
608 # Adjust for bound methods
609 folding_args = ((fnty.this,) + tuple(self.args)
610 if isinstance(fnty, types.BoundFunction)
611 else self.args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typeinfer.py:1577, in TypeInferer.resolve_call(self, fnty, pos_args, kw_args)
1574 return sig
1575 else:
1576 # Normal non-recursive call
-> 1577 return self.context.resolve_function_type(fnty, pos_args, kw_args)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:196, in BaseContext.resolve_function_type(self, func, args, kws)
194 # Prefer user definition first
195 try:
--> 196 res = self._resolve_user_function_type(func, args, kws)
197 except errors.TypingError as e:
198 # Capture any typing error
199 last_exception = e

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/typing/context.py:248, in BaseContext._resolve_user_function_type(self, func, args, kws, literals)
244 return self.resolve_function_type(func_type, args, kws)
246 if isinstance(func, types.Callable):
247 # XXX fold this into the call attribute logic?
--> 248 return func.get_call_type(self, args, kws)

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/types/functions.py:541, in Dispatcher.get_call_type(self, context, args, kws)
534 def get_call_type(self, context, args, kws):
535 """
536 Resolve a call to this dispatcher using the given argument types.
537 A signature returned and it is ensured that a compiled specialization
538 is available for it.
539 """
540 template, pysig, args, kws =
--> 541 self.dispatcher.get_call_template(args, kws)
542 sig = template(context).apply(args, kws)
543 if sig:

File ~/opt/anaconda3/lib/python3.9/site-packages/numba/core/dispatcher.py:318, in _DispatcherBase.get_call_template(self, args, kws)
316 # Ensure an overload is available
317 if self._can_compile:
--> 318 self.compile(tuple(args))
320 # Create function type for typing
321 func_name = self.py_func.name