File size: 4,491 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """Generate synthetic scRNA-seq data from known kinetic parameters + NB noise."""
from __future__ import annotations
import numpy as np
from anndata import AnnData
def generate_kinetic_data(
n_cells: int = 3000,
n_genes: int = 200,
n_cell_types: int = 5,
dispersion: float = 10.0,
sparsity: float = 0.3,
seed: int = 0,
) -> tuple[AnnData, dict[str, np.ndarray]]:
"""Generate (u, s) counts from a kinetic model with NB observation noise.
The generative process:
1. Sample latent factors z_T, z_PT per cell (different means per type).
2. Derive alpha = softplus(W_alpha @ z_T), gamma = softplus(W_gamma @ z_PT).
3. Set beta as gene-specific constants.
4. Compute mu_u ∝ alpha/beta, mu_s ∝ alpha/gamma, scaled by library size.
5. Draw counts from NB(mu, theta=dispersion).
6. Apply zero-inflation (dropout) at rate ``sparsity``.
Parameters
----------
n_cells
Number of cells.
n_genes
Number of genes.
n_cell_types
Number of simulated cell types.
dispersion
NB inverse dispersion (higher = less noise).
sparsity
Fraction of zeros injected (dropout).
seed
Random seed.
Returns
-------
adata : AnnData
With layers ``'spliced'``, ``'unspliced'``, and obs ``'cell_type'``.
truth : dict
Ground-truth arrays: ``alpha``, ``gamma``, ``beta``, ``z_T``, ``z_PT``.
"""
rng = np.random.RandomState(seed)
d_latent = 10 # latent dimension
# Cell-type assignments
cell_types = rng.choice(n_cell_types, size=n_cells)
# Cell-type-specific latent means
type_means_T = rng.randn(n_cell_types, d_latent).astype(np.float32)
type_means_PT = rng.randn(n_cell_types, d_latent).astype(np.float32)
z_T = type_means_T[cell_types] + 0.3 * rng.randn(n_cells, d_latent).astype(
np.float32
)
z_PT = type_means_PT[cell_types] + 0.3 * rng.randn(n_cells, d_latent).astype(
np.float32
)
# Decoder weights (fixed ground truth)
W_alpha = rng.randn(d_latent, n_genes).astype(np.float32) * 0.5
W_gamma = rng.randn(d_latent, n_genes).astype(np.float32) * 0.5
# Kinetic parameters
alpha = _softplus(z_T @ W_alpha) # (n_cells, n_genes)
gamma = _softplus(z_PT @ W_gamma) # (n_cells, n_genes)
beta = np.exp(rng.randn(n_genes).astype(np.float32) * 0.5 + 1.0) # gene-specific
# Expected counts (proportional)
eps = 1e-8
mu_u_raw = alpha / (beta[np.newaxis, :] + eps)
mu_s_raw = alpha / (gamma + eps)
# Library sizes
l_u = rng.lognormal(mean=8.0, sigma=0.5, size=n_cells).astype(np.float32)
l_s = rng.lognormal(mean=9.0, sigma=0.5, size=n_cells).astype(np.float32)
# Normalize to proportions then scale by library size
mu_u = (mu_u_raw / (mu_u_raw.sum(axis=1, keepdims=True) + eps)) * l_u[:, None]
mu_s = (mu_s_raw / (mu_s_raw.sum(axis=1, keepdims=True) + eps)) * l_s[:, None]
# NB sampling
u_counts = _sample_nb(mu_u, dispersion, rng)
s_counts = _sample_nb(mu_s, dispersion, rng)
# Dropout
if sparsity > 0:
mask_u = rng.rand(n_cells, n_genes) > sparsity
mask_s = rng.rand(n_cells, n_genes) > sparsity
u_counts = u_counts * mask_u
s_counts = s_counts * mask_s
adata = AnnData(
X=s_counts.astype(np.float32),
layers={
"spliced": s_counts.astype(np.float32),
"unspliced": u_counts.astype(np.float32),
},
)
adata.obs_names = [f"cell_{i}" for i in range(n_cells)]
adata.var_names = [f"gene_{i}" for i in range(n_genes)]
adata.obs["cell_type"] = [f"type_{t}" for t in cell_types]
adata.obs["cell_type"] = adata.obs["cell_type"].astype("category")
truth = {
"alpha": alpha.astype(np.float32),
"gamma": gamma.astype(np.float32),
"beta": beta.astype(np.float32),
"z_T": z_T,
"z_PT": z_PT,
}
return adata, truth
def _softplus(x: np.ndarray) -> np.ndarray:
"""Numerically stable softplus."""
return np.where(x > 20, x, np.log1p(np.exp(np.clip(x, -20, 20))))
def _sample_nb(
mu: np.ndarray, theta: float, rng: np.random.RandomState
) -> np.ndarray:
"""Sample from NB(mu, theta) using gamma-Poisson mixture."""
mu = np.clip(mu, 1e-8, None)
# Shape-rate parameterization: shape=theta, rate=theta/mu
p = theta / (theta + mu)
counts = rng.negative_binomial(theta, p)
return counts.astype(np.float32)
|