File size: 6,268 Bytes
07fcdfe | 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """Highly variable gene (HVG) selection utilities."""
import re
from typing import List, Optional, Set, Tuple
import numpy as np
def parse_perturbation_targets(name: str, gene_names: Optional[List[str]] = None) -> List[str]:
"""Parse perturbation name into target gene symbols.
Uses longest-match greedy parsing against the known gene dictionary to
correctly handle compound gene names like ``SET_KLF1`` (two genes
``SET`` and ``KLF1`` separated by ``_``, the standard delimiter in
Norman2019 and similar screens).
If *gene_names* is provided, the parser greedily matches the longest
known gene name at each position. Otherwise it falls back to simple
splitting on ``_`` and ``+``.
Parameters
----------
name : str
Raw perturbation label, e.g. ``KLF1``, ``KLF1+CEBPE``, ``SET_KLF1``.
gene_names : list of str, optional
Known gene symbols sorted by the dataset's gene dictionary.
When provided, longest-match parsing is used for correctness.
Returns
-------
list of str
Non-empty target gene symbols.
"""
name = str(name)
if gene_names is not None:
# Longest-match greedy parsing: at each position, find the longest
# gene name that matches the remaining string prefix.
sorted_genes = sorted(gene_names, key=len, reverse=True)
result = []
remaining = name
while remaining:
matched = False
for g in sorted_genes:
if remaining.startswith(g):
result.append(g)
remaining = remaining[len(g):]
matched = True
break
if not matched:
# Skip unrecognized character (separator or noise)
remaining = remaining[1:]
return result
else:
parts = re.split(r"[_+]", name)
return [p.strip() for p in parts if p.strip()]
def _is_single_gene_perturbation(name: str) -> bool:
"""Return True if the perturbation name targets exactly one gene.
Uses the same parsing logic as :func:`parse_perturbation_targets` so that
training and evaluation agree on which names are multi-gene. Names like
``SET_KLF1`` are parsed as ``["SET", "KLF1"]`` → multi-gene.
"""
return len(parse_perturbation_targets(name)) == 1
def _all_target_genes(pert_names: List[str]) -> dict:
"""Map each gene to the set of condition indices that target it.
Parameters
----------
pert_names : list of perturbation name strings
Returns
-------
dict : gene_str → set of condition indices
"""
gene_to_conds: dict = {}
for i, name in enumerate(pert_names):
for g in parse_perturbation_targets(name):
gene_to_conds.setdefault(g, set()).add(i)
return gene_to_conds
def gene_disjoint_split(
pert_names: List[str],
val_fraction: float = 0.15,
seed: int = 42,
) -> Tuple[List[int], List[int]]:
"""Split condition indices so that val perturbation genes never appear in train.
Strategy
--------
1. Collect all genes that appear **exclusively** in single-gene conditions.
(Genes that only appear in multi-gene conditions cannot be held out
without also removing all combo conditions containing them.)
2. Randomly select ``val_fraction`` of those exclusive single-gene symbols
as the val-gene set.
3. Every condition targeting *any* val-gene → validation set.
All other conditions → training set.
This guarantees zero gene overlap between train and val *for genes that
appear in single-gene conditions*. Genes that only appear in combos may
appear in both splits (since removing all their combos would discard too
much data), but this is a known and acceptable limitation.
Parameters
----------
pert_names : list of perturbation names (one per condition), e.g.
['KLF1', 'CEBPE', 'KLF1+CEBPE', 'SET_KLF1', ...]
val_fraction : fraction of exclusive single-gene symbols to hold out
seed : random seed for gene selection
Returns
-------
train_idx : list[int] condition indices for training
val_idx : list[int] condition indices for validation
"""
# Build gene → condition mapping
gene_to_conds = _all_target_genes(pert_names)
# Genes that appear ONLY in single-gene conditions (eligible for hold-out)
exclusive_single_genes = sorted({
g for g, conds in gene_to_conds.items()
if all(_is_single_gene_perturbation(pert_names[i]) for i in conds)
})
rng = np.random.default_rng(seed)
n_val = max(1, int(len(exclusive_single_genes) * val_fraction))
val_genes = set(rng.choice(exclusive_single_genes, n_val, replace=False).tolist())
train_idx, val_idx = [], []
for i, name in enumerate(pert_names):
genes = set(parse_perturbation_targets(name))
if any(g in val_genes for g in genes):
val_idx.append(i)
else:
train_idx.append(i)
return train_idx, val_idx
def select_hvg_by_variance(
X: np.ndarray,
n_genes: int = 2000,
) -> np.ndarray:
"""Select top-n genes by across-cell variance.
Parameters
----------
X : [n_cells, n_genes] float array (log-normalised)
n_genes : number of HVG to keep
Returns
-------
hvg_idx : [n_genes] integer indices into gene axis, sorted by variance desc
"""
var = X.var(axis=0)
n_genes = min(n_genes, X.shape[1])
idx = np.argsort(var)[::-1][:n_genes]
return idx.copy()
def normalize_counts(
X: np.ndarray,
target_sum: float = 1e4,
log1p: bool = True,
) -> np.ndarray:
"""Library-size normalise and optionally log1p-transform raw count matrix.
Parameters
----------
X : [n_cells, n_genes] raw count matrix (dense float32)
target_sum : scale each cell to this total count
log1p : apply log1p after normalisation
Returns
-------
X_norm : [n_cells, n_genes] float32
"""
row_sums = X.sum(axis=1, keepdims=True).clip(min=1.0)
X_norm = (X / row_sums) * target_sum
if log1p:
X_norm = np.log1p(X_norm)
return X_norm.astype(np.float32)
|