File size: 2,853 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """kNN Gaussian kernel smoothing of expression layers."""
from __future__ import annotations
import numpy as np
from anndata import AnnData
from scipy.sparse import issparse
from .._constants import (
UNSPLICED,
SPLICED,
SMOOTHED_UNSPLICED,
SMOOTHED_SPLICED,
DEFAULT_BANDWIDTH,
DEFAULT_GENE_BATCH_SIZE,
)
from .._utils import log_params, to_dense_float32
from .._numba_kernels import _smooth_kernel, _compute_adaptive_bandwidths
def smooth_layers(
adata: AnnData,
bandwidth: str | float = DEFAULT_BANDWIDTH,
gene_batch_size: int = DEFAULT_GENE_BATCH_SIZE,
layers: tuple[str, str] = (UNSPLICED, SPLICED),
out_layers: tuple[str, str] = (SMOOTHED_UNSPLICED, SMOOTHED_SPLICED),
) -> None:
"""Smooth unspliced and spliced layers using kNN Gaussian kernel.
Uses the precomputed kNN graph in ``adata.obsp['distances']``.
Gene-batched for memory control on large datasets.
Parameters
----------
adata
Annotated data matrix with precomputed neighbor graph.
bandwidth
``'adaptive'`` (default) uses the median neighbor distance per
cell. A float value sets a fixed bandwidth for all cells.
gene_batch_size
Number of genes to process per batch.
layers
Input layer names ``(unspliced, spliced)``.
out_layers
Output layer names ``(Mu, Ms)``.
"""
if "distances" not in adata.obsp:
raise ValueError(
"Neighbor graph not found. Run scptr.pp.neighbors() first."
)
dist_mat = adata.obsp["distances"]
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)
# Compute bandwidths
if bandwidth == "adaptive":
bandwidths = _compute_adaptive_bandwidths(distances, indptr)
else:
bandwidths = np.full(adata.n_obs, float(bandwidth), dtype=np.float32)
n_obs, n_vars = adata.n_obs, adata.n_vars
for in_layer, out_layer in zip(layers, out_layers):
data = to_dense_float32(adata.layers[in_layer])
result = np.empty_like(data)
# Process in gene batches
for start in range(0, n_vars, gene_batch_size):
end = min(start + gene_batch_size, n_vars)
batch_in = np.ascontiguousarray(data[:, start:end])
batch_out = np.empty_like(batch_in)
_smooth_kernel(
batch_in, indices, indptr, distances, bandwidths, batch_out
)
result[:, start:end] = batch_out
adata.layers[out_layer] = result
log_params(adata, "smooth_layers", {
"bandwidth": bandwidth,
"gene_batch_size": gene_batch_size,
})
|