| """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: |
| |
| |
| 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: |
| |
| 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 |
| """ |
| |
| gene_to_conds = _all_target_genes(pert_names) |
|
|
| |
| 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) |
|
|
