| """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) |
|
|
| |
| 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, |
| }) |
|
|
