File size: 1,963 Bytes
4e2940e | 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 | """Initialization and utility helpers for DeepPTR."""
from __future__ import annotations
from typing import TYPE_CHECKING
import numpy as np
import torch
from torch import nn
if TYPE_CHECKING:
from anndata import AnnData
def init_weights(module: nn.Module) -> None:
"""Xavier-uniform initialization for linear layers."""
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
def beta_from_adata(adata: AnnData) -> torch.Tensor:
"""Extract analytical beta estimates from *adata* for warm-starting.
Falls back to ones if ``adata.var['beta']`` is absent.
Returns
-------
torch.Tensor
Shape ``(n_genes,)``, dtype float32.
"""
if "beta" in adata.var.columns:
beta = adata.var["beta"].values.astype(np.float32)
beta = np.clip(beta, 1e-4, None)
return torch.from_numpy(np.log(beta))
return torch.zeros(adata.n_vars, dtype=torch.float32)
def get_library_sizes(adata: AnnData) -> tuple[np.ndarray, np.ndarray]:
"""Compute per-cell library sizes for unspliced and spliced layers.
Uses raw integer counts stored in ``adata.layers``.
Returns
-------
(l_u, l_s) : tuple of np.ndarray
Each shape ``(n_obs,)``, float32.
"""
from scipy.sparse import issparse
for layer in ("spliced", "unspliced"):
if layer not in adata.layers:
raise KeyError(f"Missing required layer: {layer}")
def _sum(mat: np.ndarray | "scipy.sparse.spmatrix") -> np.ndarray:
if issparse(mat):
return np.asarray(mat.sum(axis=1)).ravel().astype(np.float32)
return np.asarray(mat.sum(axis=1)).ravel().astype(np.float32)
l_s = _sum(adata.layers["spliced"])
l_u = _sum(adata.layers["unspliced"])
# Avoid zero library sizes
l_s = np.clip(l_s, 1.0, None)
l_u = np.clip(l_u, 1.0, None)
return l_u, l_s
|