File size: 1,989 Bytes
925ee3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """Post-transcriptional velocity computation."""
from __future__ import annotations
import numpy as np
from anndata import AnnData
from scipy.sparse import issparse
from .._constants import GAMMA, PT_VELOCITY
from .._utils import get_layer, require_layers, log_params
from .._numba_kernels import _velocity_kernel, _compute_adaptive_bandwidths
def pt_velocity(
adata: AnnData,
use_graph: str = "gamma",
) -> None:
"""Compute post-transcriptional velocity from gamma gradients.
For each cell, the velocity is the weighted mean difference in
gamma from its neighbors:
``v[i,g] = sum_j(w_ij * (gamma[j,g] - gamma[i,g]))``
Parameters
----------
adata
Annotated data matrix with ``gamma`` layer and neighbor graph.
use_graph
Which neighbor graph to use: ``'gamma'`` for gamma-space graph
(from ``pt_states``), ``'expression'`` for expression-space graph.
"""
require_layers(adata, GAMMA)
gamma = get_layer(adata, GAMMA).astype(np.float32)
# Get the appropriate distance matrix
if use_graph == "gamma" and "gamma_distances" in adata.obsp:
dist_mat = adata.obsp["gamma_distances"]
elif "distances" in adata.obsp:
dist_mat = adata.obsp["distances"]
else:
raise ValueError(
"No neighbor graph found. Run scptr.pp.neighbors() or "
"scptr.tl.pt_states() first."
)
if not issparse(dist_mat):
from scipy.sparse import csr_matrix
dist_mat = csr_matrix(dist_mat)
dist_mat = dist_mat.tocsr()
indices = dist_mat.indices.astype(np.int32)
indptr = dist_mat.indptr.astype(np.int32)
distances = dist_mat.data.astype(np.float32)
bandwidths = _compute_adaptive_bandwidths(distances, indptr)
out = np.empty_like(gamma)
_velocity_kernel(gamma, indices, indptr, distances, bandwidths, out)
adata.layers[PT_VELOCITY] = out
log_params(adata, "pt_velocity", {
"use_graph": use_graph,
})
|