""" Module: preprocess.py This module provides a preprocessing pipeline for single-cell RNA sequencing (scRNA-seq) data stored in AnnData format. It includes functions for loading data, filtering cells and genes, normalizing and scaling data, and saving processed results. The pipeline is designed to be configurable via hyperparameters and supports various preprocessing steps such as mitochondrial gene filtering, highly variable gene selection, and log transformation. Main Features: - Load and preprocess scRNA-seq data in AnnData format. - Filter cells and genes based on various criteria. - Normalize, scale, and log-transform data. - Save processed data and metadata to disk. - Configurable via JSON-based hyperparameters. Dependencies: - anndata, numpy, pandas, scanpy, scipy, sklearn Usage: - Run this script as a standalone program with a configuration file specifying the hyperparameters. - Import the `preprocess` function and call it with the data path, metadata path, and hyperparameters. """ import gc import json import os import warnings from argparse import ArgumentParser from typing import Sequence, Optional, Union from pathlib import Path import anndata as ad import numpy as np import pandas as pd import scanpy as sc from anndata import ImplicitModificationWarning import scipy.sparse as sp from scipy.sparse import csr_matrix, issparse from sklearn.utils import sparsefuncs, sparsefuncs_fast from teddy.data_processing.utils.gene_mapping.gene_mapper import ( map_mouse_human, map_mouse_human2, ) # --- 1. Reference list of the 37 human mitochondrial genes (Ensembl IDs) ----- _HUMAN_MITO_ENSEMBL= { "ENSG00000211459", "ENSG00000210082", # rRNAs # tRNAs (22) "ENSG00000210049", "ENSG00000210077", "ENSG00000209082", "ENSG00000210100", "ENSG00000210107", "ENSG00000210112", "ENSG00000210119", "ENSG00000210122", "ENSG00000210116", "ENSG00000210117", "ENSG00000210118", "ENSG00000210124", "ENSG00000210126", "ENSG00000210134", "ENSG00000210135", "ENSG00000210142", "ENSG00000210144", "ENSG00000210148", "ENSG00000210150", "ENSG00000210155", "ENSG00000210196", "ENSG00000210151", # protein-coding (13) "ENSG00000198888", "ENSG00000198763", "ENSG00000198840", "ENSG00000198886", "ENSG00000212907", "ENSG00000198786", "ENSG00000198695", "ENSG00000198804", "ENSG00000198712", "ENSG00000198938", "ENSG00000198899", "ENSG00000228253", "ENSG00000198727", } _HUMAN_MITO_SYMBOLS = { "MT-RNR1", "MT-RNR2", "MT-TF", "MT-TV", "MT-TL1", "MT-TI", "MT-TQ", "MT-TM", "MT-TW", "MT-TA", "MT-TN", "MT-TC", "MT-TY", "MT-TD", "MT-TK", "MT-TG", "MT-TR", "MT-TH", "MT-TS2", "MT-TL2", "MT-TT", "MT-TE", "MT-TP", "MT-TS1", "MT-ND1", "MT-ND2", "MT-ND3", "MT-ND4", "MT-ND4L", "MT-ND5", "MT-ND6", "MT-CO1", "MT-CO2", "MT-CO3", "MT-ATP6", "MT-ATP8", "MT-CYB", } def load_data_and_metadata(data_path: str, metadata_path: str): """ Load an AnnData h5ad file (data_processing) and a JSON file (metadata). """ data = ad.read_h5ad(data_path) with open(metadata_path, "r") as f: metadata = json.load(f) return data, metadata def set_raw_if_necessary(data: ad.AnnData): """ If data_processing.raw is None, checks if data_processing.X is integer for ~64 cells. If so, set data_processing.raw = data_processing. Otherwise return None (skip). """ if data.raw is not None: return data # Already has raw # If there is a 'counts' layer if 'counts' in data.layers: X = data.layers['counts'] # convert only 64 rows instead of converting the whole thing if isinstance(X, np.ndarray): X_sample = X[:64] elif issparse(X): X_sample = X[:64].toarray() # Check first 64 rows for integrality if np.all(np.equal(np.mod(X_sample, 1), 0)): data.raw = ad.AnnData(X = data.layers['counts'], var = data.var.copy()) return data # If above steps fail, check that data.X has raw counts already X = data.X # convert only 64 rows instead of converting the whole thing if isinstance(X, np.ndarray): X_sample = X[:64] elif issparse(X): X_sample = X[:64].toarray() # Check first 64 rows for integrality if np.all(np.equal(np.mod(X_sample, 1), 0)): data.raw = data return data else: print("No integer-valued matrix found") return None def initialize_processed_layer(data: ad.AnnData): """ If 'processed' layer is missing, copy from data_processing.raw.X """ if "processed" not in data.layers: data.layers["processed"] = data.raw.X.astype("float32") return data # Replacing inline code with a small helper: # (we simply inline the code from the original snippet) # You can also fully factor it out for clarity: # --------------------------------------------------- # Actually let's define that properly here to keep it consistent: def filter_reference_id(data: ad.AnnData, hyperparameters: dict): human_map = pd.read_csv("teddy/data_processing/utils/gene_mapping/data/human_mapping.txt", sep="\t") mouse_map = pd.read_csv("teddy/data_processing/utils/gene_mapping/data/2407_mouse_gene_mapping.txt", sep="\t") orthologs = pd.read_csv( "teddy/data_processing/utils/gene_mapping/data/mouse_to_human_orthologs.one2one.txt", sep="\t" ) if hyperparameters.get("mouse_nonorthologs", False): reference_id = map_mouse_human2( data_frame=data.var, query_column=None, human_map_db=human_map, mouse_map_db=mouse_map, orthology_db=orthologs, )["reference_id"] else: reference_id = map_mouse_human( data_frame=data.var, query_column=None, human_map_db=human_map, mouse_map_db=mouse_map, orthology_db=orthologs, )["reference_id"] valid_mask = reference_id != "" data = data[:, valid_mask].copy() reference_id = reference_id[valid_mask].reset_index(drop=True) if not isinstance(data.layers["processed"], np.ndarray): corrected = data.layers["processed"].toarray() else: corrected = data.layers["processed"] unique_ids = reference_id.unique() vars_to_keep = [] for rid in unique_ids: repeated_idx = np.where(reference_id == rid)[0] vars_to_keep.append(repeated_idx[0]) if len(repeated_idx) > 1: corrected[:, repeated_idx[0]] = corrected[:, repeated_idx].max(axis=1) vars_to_keep = sorted(vars_to_keep) corrected = corrected[:, vars_to_keep] data = data[:, vars_to_keep] with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=ImplicitModificationWarning) data.layers["processed"] = csr_matrix(corrected) data.var["reference_id"] = list(reference_id[vars_to_keep]) gc.collect() return data # End of inline helper # --------------------------------------------------- def remove_assays(data: ad.AnnData, assays_to_remove: list): """ Removes observations from specified 'assay' categories if 'assay' is in data_processing.obs. """ data = data[~data.obs.assay.isin(assays_to_remove)].copy() gc.collect() return data def filter_cells_by_gene_counts(data: ad.AnnData, min_count: int): """ Removes cells (observations) whose total gene counts < min_count. """ mask = sc.pp.filter_cells(data.layers["processed"], min_counts=min_count)[0] data = data[np.where(mask)].copy() del mask gc.collect() return data def filter_cells_by_mitochondrial_fraction(data: ad.AnnData, max_mito_prop: float): """ Remove low-quality cells whose mitochondrial read fraction exceeds *max_fraction*. DO NOT RUN THIS IN ANY PREPROCESSING PIPELINE UNTIL YOU HAVE SET RAW COUNTS Parameters ---------- data `AnnData` object containing counts. Works with dense or sparse matrices. max_mito_prop Threshold above which cells are discarded. Returns ------- AnnData A **copy** of `data` with poor-quality cells removed and two new columns added to ``.obs``: - **mito_prop** – per-cell mitochondrial fraction - **poor_quality_mito** – boolean flag marking dropped cells """ # We can safely assume that counts live in data.X because we set those # prior to running this step in the preprocess function. counts = data.X var_index = data.var_names if var_index[0].startswith("ENSG"): ref = _HUMAN_MITO_ENSEMBL else: ref = _HUMAN_MITO_SYMBOLS mito_idx = np.flatnonzero(var_index.isin(ref)) if mito_idx.size == 0: _logger.info("No mitochondrial genes found, returning data") return data if sp.issparse(counts): total = counts.sum(axis=1).A1 mito = counts[:, mito_idx].sum(axis=1).A1 else: total = counts.sum(axis=1) mito = counts[:, mito_idx].sum(axis=1) mito_prop = mito / np.maximum(total, 1) data.obs["mito_prop"] = mito_prop data.obs["poor_quality_mito"] = mito_prop > max_mito_prop filtered = data[~data.obs["poor_quality_mito"]].copy() gc.collect() return filtered def filter_highly_variable_genes(data: ad.AnnData, method: str): """ Filter genes to those that are highly variable using scanpy. method must be "seurat_v3" or "cell_ranger". """ if "highly_variable" in data.var: data = data[:, data.var["highly_variable"]] else: sc.pp.highly_variable_genes(data, flavor=method, n_top_genes=10000) gc.collect() return data def normalize_data_inplace(matrix_csr: csr_matrix, norm_value: float): """ In-place row normalization + scale. matrix_csr must be a CSR matrix. """ # In-place row normalize (L1) sparsefuncs_fast.inplace_csr_row_normalize_l1(matrix_csr) # Multiply each row by norm_value scale_factors = np.array([norm_value] * matrix_csr.shape[0]) sparsefuncs.inplace_row_scale(matrix_csr, scale_factors) gc.collect() def scale_columns_by_median_dict(layer: csr_matrix, data: ad.AnnData, median_dict_path: str, median_column: str): """ Read a JSON median_dict, scale columns by 1/median. The lookup key is either data_processing.var.index or data_processing.var[median_column]. """ with open(median_dict_path) as f: median_dict = json.load(f) if median_column == "index": median_var = data.var.index else: median_var = data.var[median_column] factors = [] for g in median_var: if g in median_dict: factors.append(1.0 / median_dict[g]) else: factors.append(1.0) factors = np.array(factors) # Apply in-place column scale sparsefuncs.inplace_csr_column_scale(layer, factors) def log_transform_layer(data: ad.AnnData, layer_name: str = "processed"): """ Apply sc.pp.log1p in place to data_processing.layers[layer_name]. """ sc.pp.log1p(data, layer=layer_name, copy=False) def compute_and_save_medians(data: ad.AnnData, data_path: str, hyperparameters: dict): """ Convert zeros to NaN, compute column medians ignoring NaN, and save results as JSON. """ with warnings.catch_warnings(): warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered") mat = data.layers["processed"].toarray() mat[mat == 0] = np.nan medians = np.nanmedian(mat, axis=0) if hyperparameters["median_column"] == "index": median_var = data.var.index.copy() if not isinstance(median_var, pd.Series): median_var = pd.Series(median_var) else: median_var = data.var[hyperparameters["median_column"]].copy() valid_idxs = np.where(~np.isnan(medians))[0] median_values = {median_var.iloc[k]: medians[k].item() for k in valid_idxs} save_path = data_path.replace(hyperparameters["load_dir"], hyperparameters["save_dir"]) save_path = save_path.replace(".h5ad", "_medians.json") with open(save_path, "w") as f: json.dump(median_values, f, indent=4) def update_metadata(metadata: dict, data: ad.AnnData, hyperparameters: dict): """ Update metadata with cell_count and track processing arguments. """ metadata["cell_count"] = data.n_obs if "processing_args" in metadata: metadata["processing_args"] = [metadata["processing_args"]] + [hyperparameters] else: # original fallback metadata["processings_args"] = [hyperparameters] return metadata def save_and_cleanup(data: ad.AnnData, metadata: dict, data_path: str, metadata_path: str, hyperparameters: dict): """ Write processed data_processing and metadata to disk, then GC cleanup. """ load_dir = hyperparameters["load_dir"] save_dir = hyperparameters["save_dir"] data_filename = os.path.basename(data_path) # e.g. "sample_data.h5ad" metadata_filename = os.path.basename(metadata_path) # e.g. "sample_data_metadata.json" save_processed_path = os.path.join(save_dir, data_filename) save_metadata_path = os.path.join(save_dir, metadata_filename) # Make sure the directories exist os.makedirs(os.path.dirname(save_processed_path), exist_ok=True) os.makedirs(os.path.dirname(save_metadata_path), exist_ok=True) if data.n_obs == 0: return None, None # Ensure relevant layers are sparse matrices if not isinstance(data.raw.X, csr_matrix): data.raw.X = csr_matrix(data.raw.X) if not isinstance(data.X, csr_matrix): data.X = csr_matrix(data.X) if "processed" in data.layers and not isinstance(data.layers["processed"], csr_matrix): data.layers["processed"] = csr_matrix(data.layers["processed"]) try: data.write_h5ad(save_processed_path, compression="gzip") except Exception: # Rare bug with categorical indexes if data.obs.index.name in data.obs.columns: del data.obs[data.obs.index.name] data.write_h5ad(save_processed_path, compression="gzip") del data gc.collect() with open(save_metadata_path, "w") as f: json.dump(metadata, f, indent=4) return True, True def preprocess(data_path: str, metadata_path: str, hyperparameters: dict): """ Original pipeline steps: 1. Load data_processing & metadata 2. Ensure data_processing.raw if counts are integer 3. Initialize 'processed' layer 4. Filter genes by reference_id 5. Remove assays 6. Filter cells (min gene counts) 7. Filter cells (max mito fraction) 8. HVG filtering 9. Normalize total 10. Median-based column scaling 11. Log transform 12. Compute medians (optional) 13. Update metadata and save """ # 1. Load data, metadata = load_data_and_metadata(data_path, metadata_path) # 2. Ensure data_processing.raw if needed data = set_raw_if_necessary(data) if data is None: return None, None # 3. Initialize 'processed' data = initialize_processed_layer(data) # Perturbseq fine-tuning pipeline # 4. Possible map/reference_id if hyperparameters["reference_id_only"]: data = filter_reference_id(data, hyperparameters) # 5. Remove assays if "assay" in data.obs and hyperparameters["remove_assays"]: data = remove_assays(data, hyperparameters["remove_assays"]) # 6. Filter cells by min gene counts if hyperparameters["min_gene_counts"]: data = filter_cells_by_gene_counts(data, hyperparameters["min_gene_counts"]) # 7. Filter cells by mitochondrial fraction if hyperparameters["max_mitochondrial_prop"]: # The "original" version *always* used feature_name, so fallback=False data = filter_cells_by_mitochondrial_fraction( data, hyperparameters["max_mitochondrial_prop"]) # 8. HVG filtering if hyperparameters["hvg_method"] in ["seurat_v3", "cell_ranger"]: data = filter_highly_variable_genes(data, hyperparameters["hvg_method"]) # 9. Normalize total (row L1 + scale) if hyperparameters["normalized_total"]: if not isinstance(data.layers["processed"], csr_matrix): data.layers["processed"] = csr_matrix(data.layers["processed"]) normalize_data_inplace(data.layers["processed"], hyperparameters["normalized_total"]) # 10. Scale columns using median_dict if hyperparameters["median_dict"]: scale_columns_by_median_dict( data.layers["processed"], data, hyperparameters["median_dict"], hyperparameters["median_column"] ) # 11. Log1p transform if hyperparameters["log1p"]: log_transform_layer(data, "processed") # 12. Possibly compute medians if hyperparameters["compute_medians"]: compute_and_save_medians(data, data_path, hyperparameters) # 13. Update metadata, save & cleanup metadata = update_metadata(metadata, data, hyperparameters) return save_and_cleanup(data, metadata, data_path, metadata_path, hyperparameters) ############################################################################### # Main block ############################################################################### if __name__ == "__main__": parser = ArgumentParser(description="Preprocess scRNA-seq data stored in AnnData format.") parser.add_argument( "--data_path", type=str, required=True, help="Path to the input .h5ad file." ) parser.add_argument( "--metadata_path", type=str, required=True, help="Path to the input metadata JSON file." ) parser.add_argument( "--config_path", type=str, required=True, help="Path to the JSON configuration file containing hyperparameters." ) args = parser.parse_args() # Load hyperparameters from JSON with open(args.config_path, "r") as f: hyperparameters = json.load(f) # Call the pipeline success, _ = preprocess( data_path=args.data_path, metadata_path=args.metadata_path, hyperparameters=hyperparameters ) if success: print("Preprocessing completed successfully.") else: print("Preprocessing returned no data (0 cells), no file saved.")