def __init__(
self,
obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
key: Optional[str] = None,
obsp_key: Optional[str] = None,
write_to_adata: bool = True,
):
if isinstance(obj, KernelExpression):
self._kernel = obj
elif isinstance(obj, (np.ndarray, spmatrix)):
self._kernel = PrecomputedKernel(obj)
elif isinstance(obj, AnnData):
if obsp_key is None:
raise ValueError(
"Please specify `obsp_key=...` when supplying an AnnData object."
)
elif obsp_key not in obj.obsp.keys():
raise KeyError(f"Key `{obsp_key!r}` not found in `adata.obsp`.")
self._kernel = PrecomputedKernel(obj.obsp[obsp_key], adata=obj)
else:
raise TypeError(
f"Expected an object of type `KernelExpression`, `numpy.ndarray`, `scipy.sparse.spmatrix` "
f"or `anndata.AnnData`, got `{type(obj).__name__!r}`."
)
if self.kernel.transition_matrix is None:
logg.debug("Computing transition matrix using default parameters")
self.kernel.compute_transition_matrix()
if write_to_adata:
self.kernel.write_to_adata(key=key)
def __init__(
self,
obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
key: Optional[str] = None,
obsp_key: Optional[str] = None,
write_to_adata: bool = True,
):
if isinstance(obj, KernelExpression):
self._kernel = obj
elif isinstance(obj, (np.ndarray, spmatrix)):
self._kernel = PrecomputedKernel(obj)
elif isinstance(obj, AnnData):
if obsp_key is None:
raise ValueError(
"Specify `obsp_key=...` when supplying an `AnnData` object."
)
elif obsp_key not in obj.obsp.keys():
raise KeyError(
f"Key `{obsp_key!r}` not found in `adata.obsp`.")
self._kernel = PrecomputedKernel(obj.obsp[obsp_key], adata=obj)
else:
raise TypeError(
f"Expected an object of type `KernelExpression`, `numpy.ndarray`, `scipy.sparse.spmatrix` "
f"or `anndata.AnnData`, got `{type(obj).__name__!r}`.")
if self.kernel._transition_matrix is None:
# access the private attribute to avoid accidentally computing the transition matrix
# in principle, it doesn't make a difference, apart from not seeing the message
logg.warning(
"Computing transition matrix using the default parameters")
self.kernel.compute_transition_matrix()
if write_to_adata:
self.kernel.write_to_adata(key=key)
def test_precomputed_no_adata(self):
pk = PrecomputedKernel(random_transition_matrix(50))
pk.write_to_adata()
assert isinstance(pk.adata, AnnData)
assert pk.adata.shape == (50, 1)
assert pk.adata.obs.shape == (50, 0)
assert pk.adata.var.shape == (1, 0)
assert "T_fwd_params" in pk.adata.uns.keys()
np.testing.assert_array_equal(pk.adata.obsp["T_fwd"].toarray(),
pk.transition_matrix.toarray())
def test_precomputed_from_kernel(self, adata: AnnData):
vk = VelocityKernel(adata).compute_transition_matrix(
mode="stochastic",
softmax_scale=4,
)
pk = PrecomputedKernel(vk)
pk.write_to_adata()
assert pk.adata is vk.adata
assert pk._origin == str(vk).strip("~<>")
assert pk.params is not vk.params
assert pk.params == vk.params
assert pk.transition_matrix is not vk.transition_matrix
np.testing.assert_array_equal(pk.transition_matrix.A,
vk.transition_matrix.A)
def test_precomputed_adata_origin(self, adata: AnnData):
vk = VelocityKernel(adata).compute_transition_matrix(mode="stochastic",
softmax_scale=4)
vk.write_to_adata("foo")
pk = PrecomputedKernel("foo", adata=adata)
assert pk._origin == "adata.obsp['foo']"
def test_precomputed_sum(self, adata: AnnData):
mat = random_transition_matrix(adata.n_obs)
pk = PrecomputedKernel(mat)
vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
expected = (0.5 * vk.transition_matrix) + (0.5 * pk.transition_matrix)
actual = (pk + vk).compute_transition_matrix()
np.testing.assert_array_almost_equal(
expected.toarray(), actual.transition_matrix.toarray())
def test_precomputed_different_adata(self, adata: AnnData):
vk = VelocityKernel(adata).compute_transition_matrix(mode="stochastic",
softmax_scale=4)
bdata = adata.copy()
pk = PrecomputedKernel(vk, adata=bdata)
assert pk.adata is adata
assert pk.adata is vk.adata
assert pk.adata is not bdata
def write(self,
fname: Union[str, Path],
ext: Optional[str] = "pickle") -> None:
"""
%(pickleable.full_desc)s
Parameters
----------
%(pickleable.parameters)s
Returns
-------
%(pickleable.returns)s
""" # noqa
fname = str(fname)
if ext is not None:
if not ext.startswith("."):
ext = "." + ext
if not fname.endswith(ext):
fname += ext
logg.debug(f"Writing to `{fname}`")
with open(fname, "wb") as fout:
if version_info[:2] > (3, 6):
pickle.dump(self, fout)
else:
# we need to use PrecomputedKernel because Python3.6 can't pickle Enums
# and they are present in VelocityKernel
logg.warning(
"Saving kernel as `cellrank.tl.kernels.PrecomputedKernel`")
orig_kernel = self.kernel
self._kernel = PrecomputedKernel(self.kernel)
try:
pickle.dump(self, fout)
except Exception as e:
raise e
finally:
self._kernel = orig_kernel
def test_precomputed_transition_matrix(self, adata: AnnData):
mat = random_transition_matrix(adata.n_obs)
pk = PrecomputedKernel(mat)
np.testing.assert_array_equal(mat, pk.transition_matrix.toarray())
def test_precomputed_adata(self, adata: AnnData):
pk = PrecomputedKernel(random_transition_matrix(adata.n_obs),
adata=adata)
assert pk.adata is adata
def test_precomputed_not_a_transition_matrix(self):
mat = random_transition_matrix(100)
mat[0, 0] = 0xDEADBEEF
with pytest.raises(ValueError):
_ = PrecomputedKernel(mat)
def test_precomputed_not_square(self):
with pytest.raises(ValueError):
_ = PrecomputedKernel(np.random.normal(size=(10, 9)))
def test_precomputed_not_array(self):
with pytest.raises(TypeError):
_ = PrecomputedKernel([[1, 0], [0, 1]])
def _initial_terminal(
adata: AnnData,
estimator: type(BaseEstimator) = GPCCA,
backward: bool = False,
mode: str = VelocityMode.DETERMINISTIC.s,
backward_mode: str = BackwardMode.TRANSPOSE.s,
n_states: Optional[int] = None,
cluster_key: Optional[str] = None,
key: Optional[str] = None,
show_plots: bool = False,
copy: bool = False,
return_estimator: bool = False,
fit_kwargs: Mapping = MappingProxyType({}),
**kwargs,
) -> Optional[Union[AnnData, BaseEstimator]]:
_check_estimator_type(estimator)
try:
kernel = PrecomputedKernel(key, adata=adata, backward=backward)
write_to_adata = False # no need to write
logg.info("Using precomputed transition matrix")
except KeyError:
# compute kernel object
kernel = transition_matrix(
adata,
backward=backward,
mode=mode,
backward_mode=backward_mode,
**kwargs,
)
write_to_adata = True
# create estimator object
mc = estimator(
kernel,
read_from_adata=False,
inplace=not copy,
key=key,
write_to_adata=write_to_adata,
)
if cluster_key is None:
_info_if_obs_keys_categorical_present(
adata,
keys=["louvain", "leiden", "clusters"],
msg_fmt=
"Found categorical observation in `adata.obs[{!r}]`. Consider specifying it as `cluster_key`.",
)
mc.fit(
n_lineages=n_states,
cluster_key=cluster_key,
compute_absorption_probabilities=False,
**fit_kwargs,
)
if show_plots:
mc.plot_spectrum(real_only=True)
if isinstance(mc, CFLARE):
mc.plot_eigendecomposition(abs_value=True,
perc=[0, 98],
use=n_states)
mc.plot_terminal_states(discrete=True, same_plot=False)
elif isinstance(mc, GPCCA):
n_states = len(mc._get(P.MACRO).cat.categories)
if n_states > 1:
mc.plot_schur()
mc.plot_terminal_states(discrete=True, same_plot=False)
if n_states > 1:
mc.plot_coarse_T()
else:
raise NotImplementedError(
f"Pipeline not implemented for `{type(mc).__name__!r}.`")
return mc.adata if copy else mc if return_estimator else None
def lineages(
adata: AnnData,
backward: bool = False,
copy: bool = False,
return_estimator: bool = False,
**kwargs,
) -> Optional[AnnData]:
"""
Compute probabilistic lineage assignment using RNA velocity.
For each cell `i` in :math:`{1, ..., N}` and %(initial_or_terminal)s state `j` in :math:`{1, ..., M}`,
the probability is computed that cell `i` is either going to %(terminal)s state `j` (``backward=False``)
or is coming from %(initial)s state `j` (``backward=True``).
This function computes the absorption probabilities of a Markov chain towards the %(initial_or_terminal) states
uncovered by :func:`cellrank.tl.initial_states` or :func:`cellrank.tl.terminal_states` using a highly efficient
implementation that scales to large cell numbers.
It's also possible to calculate mean and variance of the time until absorption for all or just a subset
of the %(initial_or_terminal)s states. This can be seen as a pseudotemporal measure, either towards any terminal
population of the state change trajectory, or towards specific ones.
Parameters
----------
%(adata)s
%(backward)s
copy
Whether to update the existing ``adata`` object or to return a copy.
return_estimator
Whether to return the estimator. Only available when ``copy=False``.
**kwargs
Keyword arguments for :meth:`cellrank.tl.estimators.BaseEstimator.compute_absorption_probabilities`.
Returns
-------
:class:`anndata.AnnData`, :class:`cellrank.tl.estimators.BaseEstimator` or :obj:`None`
Depending on ``copy`` and ``return_estimator``, either updates the existing ``adata`` object,
returns its copy or returns the estimator.
"""
if backward:
lin_key = AbsProbKey.BACKWARD
fs_key = TermStatesKey.BACKWARD
fs_key_pretty = TerminalStatesPlot.BACKWARD
else:
lin_key = AbsProbKey.FORWARD
fs_key = TermStatesKey.FORWARD
fs_key_pretty = TerminalStatesPlot.FORWARD
try:
pk = PrecomputedKernel(adata=adata, backward=backward)
except KeyError as e:
raise RuntimeError(
f"Compute transition matrix first as `cellrank.tl.transition_matrix(..., backward={backward})`."
) from e
start = logg.info(
f"Computing lineage probabilities towards {fs_key_pretty.s}")
mc = GPCCA(
pk, read_from_adata=True, inplace=not copy
) # GPCCA is more general than CFLARE, in terms of what is saves
if mc._get(P.TERM) is None:
raise RuntimeError(
f"Compute the states first as `cellrank.tl.{fs_key.s}(..., backward={backward})`."
)
# compute the absorption probabilities
mc.compute_absorption_probabilities(**kwargs)
logg.info(f"Adding lineages to `adata.obsm[{lin_key.s!r}]`\n Finish",
time=start)
return mc.adata if copy else mc if return_estimator else None
def test_precomputed_from_kernel_no_transition(self, adata: AnnData):
vk = VelocityKernel(adata)
with pytest.raises(ValueError):
PrecomputedKernel(vk)