Source code for cellrank.tl._transition_matrix

# -*- coding: utf-8 -*-
"""Transition matrix module."""

from typing import TypeVar, Iterable, Optional

from cellrank import logging as logg
from cellrank.ul._docs import d, inject_docs
from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
from cellrank.tl.kernels._base_kernel import KernelExpression
from cellrank.tl.kernels._velocity_kernel import BackwardMode, VelocityMode
from cellrank.tl.kernels._velocity_schemes import Scheme

AnnData = TypeVar("AnnData")


[docs]@inject_docs(m=VelocityMode, b=BackwardMode, s=Scheme) # don't swap the order @d.dedent def transition_matrix( adata: AnnData, backward: bool = False, vkey: str = "velocity", xkey: str = "Ms", gene_subset: Optional[Iterable] = None, mode: str = VelocityMode.DETERMINISTIC.s, backward_mode: str = BackwardMode.TRANSPOSE.s, scheme: str = Scheme.CORRELATION.s, softmax_scale: Optional[float] = None, weight_connectivities: Optional[float] = 0.2, density_normalize: bool = True, key: Optional[str] = None, **kwargs, ) -> KernelExpression: """ Compute a transition matrix based on a combination of RNA Velocity and transcriptomic similarity. To learn more about the way in which the transition matrices are computed, see :class:`cellrank.tl.kernels.VelocityKernel` for the velocity-based transition matrix and :class:`cellrank.tl.kernels.ConnectivityKernel` for the transcriptomic-similarity-based transition matrix. Parameters ---------- %(adata)s %(backward)s vkey Key from ``adata.layers`` to access the velocities. xkey Key in ``adata.layers`` where expected gene expression counts are stored. gene_subset List of genes to be used to compute transition probabilities. By default, genes from ``adata.var['velocity_genes']`` are used. %(velocity_mode)s %(velocity_backward_mode_high_lvl)s %(velocity_scheme)s %(softmax_scale)s weight_connectivities Weight given to transcriptomic similarities as opposed to velocities. Must be in `[0, 1]`. density_normalize Whether to use density correction when computing the transition probabilities based on connectivities. Density correction is done as by [Haghverdi16]_. %(write_to_adata.parameters)s **kwargs Keyword arguments for :meth:`cellrank.tl.kernels.VelocityKernel.compute_transition_matrix`. Returns ------- :class:`cellrank.tl.KernelExpression` A kernel expression object containing the computed transition matrix. %(write_to_adata)s """ # initialise the velocity kernel and compute transition matrix vk = VelocityKernel( adata, backward=backward, vkey=vkey, xkey=xkey, gene_subset=gene_subset ) vk.compute_transition_matrix( softmax_scale=softmax_scale, mode=mode, backward_mode=backward_mode, scheme=scheme, **kwargs, ) if weight_connectivities is not None: if 0 < weight_connectivities < 1: logg.info( f"Using a connectivity kernel with weight `{weight_connectivities}`" ) ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix( density_normalize=density_normalize ) final = ( (1 - weight_connectivities) * vk + weight_connectivities * ck ).compute_transition_matrix() elif weight_connectivities == 0: final = vk elif weight_connectivities == 1: final = ConnectivityKernel( adata, backward=backward ).compute_transition_matrix(density_normalize=density_normalize) else: raise ValueError( f"Parameter `weight_connectivities` must be in range `[0, 1]`, found `{weight_connectivities}`." ) else: final = vk final.write_to_adata(key=key) return final