| """
|
| Adaptation code of the BioMod model on the zero-shot batch correction task.
|
|
|
| The current script mainly supports:
|
| 1. Use the BioMod model (encoder-only) to extract cell embeddings.
|
| 2. Use mean pooling of module encodings as cell embedding.
|
| 3. Supports evaluation based on scIB indicators consistent with the scGPT tutorial.
|
| 4. Supports Lung data set preset and retains all genes by default without going through the HVG 3000 process.
|
| """
|
|
|
| import argparse
|
| import importlib.util
|
| import os
|
| import pickle
|
| import sys
|
| import tempfile
|
| from pathlib import Path
|
|
|
| REPO_ROOT = Path(__file__).resolve().parents[2]
|
| ASSET_DIR = REPO_ROOT / "model_dependencies"
|
| DOWNSTREAM_DATA_DIR = REPO_ROOT / "data" / "downstream" / "final_data"
|
| WEIGHTS_DIR = REPO_ROOT / "weight"
|
| from copy import deepcopy
|
| from typing import Dict, List, Optional, Sequence, Set, Union
|
|
|
| os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
|
| os.environ.setdefault("NUMBA_CACHE_DIR", tempfile.mkdtemp(prefix="numba_cache_"))
|
|
|
| import numpy as np
|
| import pandas as pd
|
| import scanpy as sc
|
| import scib
|
| import torch
|
| from scanpy import AnnData
|
| from tqdm import tqdm
|
|
|
|
|
|
|
|
|
| DATASET_PRESETS = {
|
| "auto": {},
|
| "lung": {
|
| "gene_col": "gene_name",
|
| "batch_key": "sample",
|
| "label_keys": ["cell_type"],
|
| },
|
| }
|
|
|
|
|
| def print_sys(message: str) -> None:
|
| print(message, flush=True, file=sys.stderr)
|
|
|
|
|
| def strip_gene_version(gene_id: str) -> str:
|
| return gene_id.split(".", 1)[0]
|
|
|
|
|
| def canonicalize_target_gene_id(
|
| raw_gene: object,
|
| symbol_to_ensembl: Dict[str, str],
|
| target_gene_set: Set[str],
|
| ) -> Optional[str]:
|
| if pd.isna(raw_gene):
|
| return None
|
|
|
| raw_gene_str = str(raw_gene).strip()
|
| if not raw_gene_str:
|
| return None
|
|
|
| direct_candidates = [raw_gene_str, strip_gene_version(raw_gene_str)]
|
| for candidate in direct_candidates:
|
| if candidate in target_gene_set:
|
| return candidate
|
|
|
| mapped_gene = symbol_to_ensembl.get(raw_gene_str)
|
| if mapped_gene is None:
|
| mapped_gene = symbol_to_ensembl.get(strip_gene_version(raw_gene_str))
|
|
|
| if mapped_gene is None or pd.isna(mapped_gene):
|
| return None
|
|
|
| mapped_gene_str = strip_gene_version(str(mapped_gene).strip())
|
| return mapped_gene_str or None
|
|
|
|
|
| def resolve_input_gene_ids(
|
| adata: AnnData,
|
| gene_col: Optional[str],
|
| symbol_to_ensembl: Dict[str, str],
|
| target_gene_set: Set[str],
|
| ) -> List[str]:
|
| if gene_col is None:
|
| raw_genes = [str(gene) for gene in adata.var.index.tolist()]
|
| source_name = "adata.var.index"
|
| else:
|
| if gene_col not in adata.var.columns:
|
| raise KeyError(f"gene_col '{gene_col}' not found in adata.var.columns")
|
| raw_genes = [str(gene) if not pd.isna(gene) else "" for gene in adata.var[gene_col].tolist()]
|
| source_name = f"adata.var['{gene_col}']"
|
|
|
| resolved_gene_ids: List[str] = []
|
| direct_hits = 0
|
| mapped_hits = 0
|
| unresolved = 0
|
|
|
| for index, raw_gene in enumerate(raw_genes):
|
| resolved_gene = canonicalize_target_gene_id(raw_gene, symbol_to_ensembl, target_gene_set)
|
| if resolved_gene is None:
|
| resolved_gene = f"unknown{index}"
|
| unresolved += 1
|
| elif strip_gene_version(raw_gene) == resolved_gene:
|
| direct_hits += 1
|
| else:
|
| mapped_hits += 1
|
| resolved_gene_ids.append(resolved_gene)
|
|
|
| print_sys(f"Using genes from {source_name}")
|
| print_sys(
|
| "Gene ID resolution summary: "
|
| f"direct_hits={direct_hits}, mapped_hits={mapped_hits}, unresolved={unresolved}"
|
| )
|
| return resolved_gene_ids
|
|
|
|
|
| def clear_neighbors(adata: AnnData) -> None:
|
| neighbors = adata.uns.pop("neighbors", None)
|
| if neighbors is not None:
|
| dist_key = neighbors.get("distances_key", "distances")
|
| conn_key = neighbors.get("connectivities_key", "connectivities")
|
| adata.obsp.pop(dist_key, None)
|
| adata.obsp.pop(conn_key, None)
|
| adata.obsp.pop("distances", None)
|
| adata.obsp.pop("connectivities", None)
|
|
|
|
|
| def prepare_eval_adata(
|
| adata: AnnData,
|
| batch_key: str,
|
| label_key: str,
|
| embedding_key: str,
|
| build_neighbors: bool = False,
|
| ) -> AnnData:
|
| if embedding_key not in adata.obsm:
|
| raise ValueError(f"Embeddings '{embedding_key}' not found in adata.obsm")
|
|
|
| if batch_key not in adata.obs.columns:
|
| raise ValueError(f"batch_key '{batch_key}' not found in adata.obs")
|
|
|
| if label_key not in adata.obs.columns:
|
| raise ValueError(f"label_key '{label_key}' not found in adata.obs")
|
|
|
| valid_mask = adata.obs[batch_key].notna() & adata.obs[label_key].notna()
|
| dropped = int((~valid_mask).sum())
|
| if dropped > 0:
|
| print_sys(
|
| f"Dropping {dropped} cells with missing batch/label values "
|
| f"for evaluation (batch_key={batch_key}, label_key={label_key})."
|
| )
|
|
|
| if int(valid_mask.sum()) == 0:
|
| raise ValueError(
|
| f"No labeled cells remain after filtering batch_key={batch_key}, label_key={label_key}"
|
| )
|
|
|
| eval_adata = adata[valid_mask].copy()
|
| eval_adata.obs[batch_key] = eval_adata.obs[batch_key].astype("category")
|
| eval_adata.obs[label_key] = eval_adata.obs[label_key].astype("category")
|
|
|
| if build_neighbors:
|
| clear_neighbors(eval_adata)
|
| sc.pp.neighbors(eval_adata, use_rep=embedding_key)
|
| return eval_adata
|
|
|
|
|
| def eval_scib_metrics_simple(
|
| adata: AnnData,
|
| batch_key: str,
|
| label_key: str,
|
| embedding_key: str,
|
| ) -> Dict[str, float]:
|
| eval_adata = prepare_eval_adata(
|
| adata,
|
| batch_key,
|
| label_key,
|
| embedding_key,
|
| build_neighbors=False,
|
| )
|
|
|
| results_dict: Dict[str, float] = {}
|
| if len(eval_adata.obs[batch_key].unique()) > 1:
|
| results_dict["ASW_label/batch"] = scib.metrics.silhouette_batch(
|
| eval_adata,
|
| batch_key,
|
| label_key,
|
| embed=embedding_key,
|
| metric="euclidean",
|
| return_all=False,
|
| verbose=False,
|
| )
|
|
|
| results_dict = {key: value for key, value in results_dict.items() if not np.isnan(value)}
|
| return results_dict
|
|
|
|
|
| def eval_scib_metrics_scgpt(
|
| adata: AnnData,
|
| batch_key: str,
|
| label_key: str,
|
| embedding_key: str,
|
| ) -> Dict[str, float]:
|
| eval_adata = prepare_eval_adata(
|
| adata,
|
| batch_key,
|
| label_key,
|
| embedding_key,
|
| build_neighbors=True,
|
| )
|
|
|
| results = scib.metrics.metrics(
|
| eval_adata,
|
| adata_int=eval_adata,
|
| batch_key=batch_key,
|
| label_key=label_key,
|
| embed=embedding_key,
|
| isolated_labels_asw_=False,
|
| silhouette_=True,
|
| hvg_score_=False,
|
| graph_conn_=True,
|
| pcr_=True,
|
| isolated_labels_f1_=False,
|
| trajectory_=False,
|
| nmi_=True,
|
| ari_=True,
|
| cell_cycle_=False,
|
| kBET_=False,
|
| ilisi_=False,
|
| clisi_=False,
|
| )
|
| result_dict = results[0].to_dict()
|
|
|
| result_dict["avg_bio"] = np.mean(
|
| [
|
| result_dict["NMI_cluster/label"],
|
| result_dict["ARI_cluster/label"],
|
| result_dict["ASW_label"],
|
| ]
|
| )
|
| result_dict["avg_batch"] = np.mean(
|
| [
|
| result_dict["graph_conn"],
|
| result_dict["ASW_label/batch"],
|
| ]
|
| )
|
|
|
| result_dict = {key: value for key, value in result_dict.items() if not np.isnan(value)}
|
| return result_dict
|
|
|
|
|
| def evaluate(
|
| adata: AnnData,
|
| batch_key: str,
|
| label_keys: Sequence[str],
|
| embedding_key: str,
|
| metrics_mode: str,
|
| res_path: str,
|
| ) -> pd.DataFrame:
|
| met_df = pd.DataFrame(columns=["metric", "label", "value"])
|
|
|
| unique_label_keys: List[str] = []
|
| for key in label_keys:
|
| if key not in unique_label_keys:
|
| unique_label_keys.append(key)
|
|
|
| if not unique_label_keys:
|
| raise ValueError("No label keys were provided for evaluation")
|
|
|
| for label_key in unique_label_keys:
|
| if metrics_mode == "scgpt":
|
| metrics = eval_scib_metrics_scgpt(
|
| adata=adata,
|
| batch_key=batch_key,
|
| label_key=label_key,
|
| embedding_key=embedding_key,
|
| )
|
| elif metrics_mode == "simple":
|
| metrics = eval_scib_metrics_simple(
|
| adata=adata,
|
| batch_key=batch_key,
|
| label_key=label_key,
|
| embedding_key=embedding_key,
|
| )
|
| else:
|
| raise ValueError(f"Unknown metrics_mode: {metrics_mode}")
|
|
|
| print_sys(
|
| "\n".join(
|
| [f"[{label_key}] {metric_name}: {metric_value:.4f}" for metric_name, metric_value in metrics.items()]
|
| )
|
| )
|
| for metric_name, metric_value in metrics.items():
|
| met_df.loc[len(met_df)] = [metric_name, label_key, metric_value]
|
|
|
| met_df.to_csv(res_path, index=False)
|
| np.save(f"{res_path[:-4]}_ce.npy", adata.obsm[embedding_key])
|
| return met_df
|
|
|
|
|
| def resolve_with_preset(
|
| cli_value: Optional[object],
|
| preset_value: Optional[object],
|
| ) -> Optional[object]:
|
| return cli_value if cli_value is not None else preset_value
|
|
|
|
|
| def resolve_batch_key(adata: AnnData, requested_batch_key: Optional[str]) -> str:
|
| if requested_batch_key is not None:
|
| if requested_batch_key not in adata.obs.columns:
|
| raise KeyError(f"batch_key '{requested_batch_key}' not found in adata.obs.columns")
|
| return requested_batch_key
|
|
|
| fallback_keys = ["batch", "str_batch", "sample", "Sample Characteristic[individual]"]
|
| for key in fallback_keys:
|
| if key in adata.obs.columns:
|
| print_sys(f"Auto-selected batch_key='{key}'")
|
| return key
|
|
|
| raise ValueError(
|
| "Unable to infer batch_key. Please pass --batch_key explicitly."
|
| )
|
|
|
|
|
| def resolve_label_keys(adata: AnnData, requested_label_keys: Optional[Sequence[str]]) -> List[str]:
|
| if requested_label_keys:
|
| missing_keys = [key for key in requested_label_keys if key not in adata.obs.columns]
|
| if missing_keys:
|
| raise KeyError(f"label_key(s) not found in adata.obs.columns: {missing_keys}")
|
| return list(requested_label_keys)
|
|
|
| fallback_keys = ["celltype", "cell_type", "celltype_original"]
|
| resolved_keys = [key for key in fallback_keys if key in adata.obs.columns]
|
| if resolved_keys:
|
| print_sys(f"Auto-selected label_keys={resolved_keys}")
|
| return resolved_keys
|
|
|
| raise ValueError(
|
| "Unable to infer label_key. Please pass --label_key explicitly."
|
| )
|
|
|
|
|
| def reindex_tensor_universal(
|
| tensor: torch.Tensor,
|
| index_positions: Sequence[int],
|
| dim: int,
|
| filler: str = "0",
|
| device: Union[str, torch.device] = "cpu",
|
| ) -> torch.Tensor:
|
| index_tensor = torch.tensor(index_positions, device=device)
|
| tensor_shape = list(tensor.shape)
|
| tensor_shape[dim] = len(index_positions)
|
| index_shape = [1] * len(tensor_shape)
|
| index_shape[dim] = len(index_positions)
|
| index_tensor = index_tensor.view(*index_shape).expand(*tensor_shape)
|
|
|
| if filler == "0":
|
| padder_shape = deepcopy(tensor_shape)
|
| padder_shape[dim] = 1
|
| padder = torch.zeros(padder_shape, dtype=tensor.dtype, device=device)
|
| elif filler == "mean":
|
| padder = tensor.mean(dim=dim, keepdim=True)
|
| else:
|
| raise ValueError("filler should be 0 or mean")
|
| expanded_tensor = torch.cat((tensor, padder), dim=dim)
|
| return torch.gather(expanded_tensor, dim, index_tensor)
|
|
|
|
|
| project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
| sys.path.insert(0, project_root)
|
| sys.path.insert(0, os.path.dirname(project_root))
|
|
|
| biomod_dir = project_root
|
|
|
| def load_biomod_components():
|
| models_path = os.path.join(biomod_dir, "models_biomod.py")
|
| if not os.path.exists(models_path):
|
| raise FileNotFoundError(f"Model file not found: {models_path}")
|
| models_spec = importlib.util.spec_from_file_location("models_biomod", models_path)
|
| models_biomod = importlib.util.module_from_spec(models_spec)
|
| models_spec.loader.exec_module(models_biomod)
|
|
|
| dataset_path = os.path.join(biomod_dir, "biomod_dataset.py")
|
| if not os.path.exists(dataset_path):
|
| raise FileNotFoundError(f"Dataset file not found: {dataset_path}")
|
| dataset_spec = importlib.util.spec_from_file_location("biomod_dataset", dataset_path)
|
| biomod_dataset = importlib.util.module_from_spec(dataset_spec)
|
| dataset_spec.loader.exec_module(biomod_dataset)
|
|
|
| return (
|
| models_biomod.BioMod,
|
| models_biomod.sanitize_biomod_config,
|
| biomod_dataset.load_biomod_data,
|
| )
|
|
|
| OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "res")
|
| os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument("--gene2vec_path", type=str, default=str(ASSET_DIR / "selected_gene2vec_27k.npy"),
|
| help="Path to gene2vec file")
|
| parser.add_argument("--scfm_hyper_params_path", type=str, required=True,
|
| help="Path to the matching BioMod model configuration pkl file")
|
| parser.add_argument("--scfm_ckpt_path", type=str, required=True,
|
| help="Path to the BioMod model weights")
|
| parser.add_argument("--biomod_data_path", "--module_data_path", type=str, default=None,
|
| help="Path to module JSON file (single file)", dest='module_data_path')
|
| parser.add_argument("--biomod_data_dir", "--module_data_dir", type=str, default=None,
|
| help="Path to directory containing module JSON files", dest='module_data_dir')
|
| parser.add_argument("--gene_mapping_path", type=str, default=str(ASSET_DIR / "gene2ens.csv"),
|
| help="Path to gene symbol to Ensembl ID mapping CSV")
|
| parser.add_argument("--scfm_genes_file", type=str, default=str(ASSET_DIR / "selected_genes_27k.txt"),
|
| help="Path to scfm genes list file")
|
| parser.add_argument("--data_path", type=str, required=True,
|
| help="Path to input h5ad file")
|
| parser.add_argument("--dataset_preset", type=str, default="auto", choices=sorted(DATASET_PRESETS.keys()),
|
| help="Dataset preset for batch/label/gene column defaults")
|
| parser.add_argument("--gene_col", type=str, default=None,
|
| help="Column in adata.var used as the input gene symbol column. "
|
| "If not set, use adata.var.index.")
|
| parser.add_argument("--batch_key", type=str, default=None,
|
| help="Column in adata.obs used as the batch label")
|
| parser.add_argument("--label_key", nargs="+", default=None,
|
| help="One or more columns in adata.obs used as evaluation labels")
|
| parser.add_argument("--output_key", type=str, default="X_sclong_module",
|
| help="Key name for storing embeddings in adata.obsm")
|
| parser.add_argument("--pool_method", type=str, default="auto", choices=["auto", "mean", "cls", "max", "first_module"],
|
| help="Cell embedding pooling. auto prefers cls_encoding when the checkpoint has a trained CLS token; "
|
| "first_module uses the first module after restoring original module order.")
|
| parser.add_argument("--n_cells", type=str, default="all",
|
| help="Number of cells to process (default: all)")
|
| parser.add_argument("--batch_size", type=int, default=4,
|
| help="Batch size for inference")
|
| parser.add_argument("--output_dir", type=str, default=OUTPUT_DIR,
|
| help="Output directory for results")
|
| parser.add_argument("--metrics_mode", type=str, default="scgpt", choices=["scgpt", "simple"],
|
| help="Evaluation metrics mode. Default uses the scGPT tutorial metric bundle.")
|
| parser.add_argument("--min_genes_per_cell", type=int, default=0,
|
| help="Optional Scanpy cell filter. Default 0 keeps all cells.")
|
| parser.add_argument("--min_cells_per_gene", type=int, default=0,
|
| help="Optional Scanpy gene filter. Default 0 keeps all genes.")
|
| parser.add_argument("--shuffle_biomod_genes", "--shuffle_module_genes", action="store_true",
|
| help="Optional downstream-only extra reshuffle on top of module_sets restored from PKL. "
|
| "Do not pass this when reproducing the pretrained backbone as-is.", dest='shuffle_module_genes')
|
| parser.add_argument("--shuffle_seed", type=int, default=2021,
|
| help="Random seed for the downstream extra reshuffle when --shuffle_biomod_genes is enabled.")
|
| parser.add_argument("--save_raw_biomod_emb", "--save_raw_module_emb", action="store_true",
|
| help="Save raw module_encodings with shape (num_cells, num_modules, embed_dim)", dest='save_raw_module_emb')
|
|
|
| args = parser.parse_args()
|
| os.makedirs(args.output_dir, exist_ok=True)
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| print_sys(f"Using device: {device}")
|
|
|
| preset_config = DATASET_PRESETS[args.dataset_preset]
|
| resolved_gene_col = resolve_with_preset(args.gene_col, preset_config.get("gene_col"))
|
| resolved_batch_key = resolve_with_preset(args.batch_key, preset_config.get("batch_key"))
|
| resolved_label_keys = resolve_with_preset(args.label_key, preset_config.get("label_keys"))
|
|
|
| print_sys(
|
| f"dataset_preset={args.dataset_preset}, "
|
| f"gene_col={resolved_gene_col}, batch_key={resolved_batch_key}, "
|
| f"label_keys={resolved_label_keys}, metrics_mode={args.metrics_mode}, "
|
| f"pool_method={args.pool_method}"
|
| )
|
|
|
| BioMod, sanitize_biomod_config, load_biomod_data = load_biomod_components()
|
|
|
| print_sys(f"Loading hyper parameters from {args.scfm_hyper_params_path}")
|
| with open(args.scfm_hyper_params_path, "rb") as handle:
|
| scfm_hyper_params = pickle.load(handle)
|
| scfm_hyper_params.pop("module_embed_chunk_padded_gene_budget", None)
|
| scfm_hyper_params = sanitize_biomod_config(
|
| scfm_hyper_params,
|
| source=args.scfm_hyper_params_path,
|
| log_fn=print_sys,
|
| )
|
|
|
| scfm_hyper_params["gene2vec_file"] = args.gene2vec_path
|
|
|
| fusion_method = scfm_hyper_params.get("fusion_method", "multiply")
|
| print_sys(f"Using fusion_method from pkl: {fusion_method}")
|
|
|
| if "module_sets" not in scfm_hyper_params or scfm_hyper_params.get("module_sets") is None:
|
| if args.module_data_path or args.module_data_dir:
|
| print_sys("Loading module data...")
|
| module_sets = load_biomod_data(
|
| module_data_path=args.module_data_path,
|
| module_data_dir=args.module_data_dir,
|
| gene_list_path=args.scfm_genes_file,
|
| gene_mapping_path=args.gene_mapping_path,
|
| verbose=True,
|
| )
|
| scfm_hyper_params["module_sets"] = module_sets
|
| else:
|
| raise ValueError("Model requires module_sets, but neither --biomod_data_path nor --biomod_data_dir was provided")
|
| else:
|
| module_sets = scfm_hyper_params["module_sets"]
|
| print_sys(f"Using module_sets from hyper parameters ({len(module_sets)} modules)")
|
|
|
| if args.shuffle_module_genes:
|
| import random
|
|
|
| print_sys(
|
| f"Applying downstream-only extra module reshuffle (seed={args.shuffle_seed}). "
|
| f"Pretraining-time module_sets from PKL will be reshuffled again."
|
| )
|
| all_gene_indices: List[int] = []
|
| module_sizes: List[int] = []
|
| for module in module_sets:
|
| all_gene_indices.extend(module)
|
| module_sizes.append(len(module))
|
|
|
| rng = random.Random(args.shuffle_seed)
|
| rng.shuffle(all_gene_indices)
|
|
|
| shuffled_module_sets = []
|
| start_idx = 0
|
| for size in module_sizes:
|
| shuffled_module_sets.append(all_gene_indices[start_idx:start_idx + size])
|
| start_idx += size
|
| module_sets = shuffled_module_sets
|
| scfm_hyper_params["module_sets"] = module_sets
|
| print_sys(f"Shuffled {len(all_gene_indices)} genes across {len(module_sets)} modules")
|
|
|
| print_sys(f"Model hyper parameters: {list(scfm_hyper_params.keys())}")
|
|
|
| pkl_use_gene2vec = scfm_hyper_params.get("use_gene2vec", True)
|
| scfm_hyper_params["use_gene2vec"] = pkl_use_gene2vec
|
| print_sys(f"use_gene2vec={pkl_use_gene2vec} (from PKL)")
|
|
|
| print_sys("Creating BioMod model...")
|
| model = BioMod(**scfm_hyper_params).to(device)
|
|
|
| print_sys(f"Loading checkpoint from {args.scfm_ckpt_path}")
|
| ckpt = torch.load(args.scfm_ckpt_path, map_location=device)
|
| model_state_dict = ckpt.get("model_state_dict", ckpt)
|
| filtered_state_dict = {
|
| key: value
|
| for key, value in model_state_dict.items()
|
| if "go_conv" not in key and "G_go" not in key.lower()
|
| }
|
| model.load_state_dict(filtered_state_dict, strict=True)
|
|
|
| for parameter in model.parameters():
|
| parameter.requires_grad = False
|
| model.eval()
|
|
|
| all_params = sum(parameter.numel() for parameter in model.parameters())
|
| trainable_params = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
|
| print_sys(f"Total model parameters: {all_params:,}")
|
| print_sys(f"Trainable parameters during zero-shot: {trainable_params:,}")
|
|
|
| gene2ens = pd.read_csv(args.gene_mapping_path, index_col=0)
|
| symbol_to_ensembl = gene2ens["Ensembl gene ID"].to_dict()
|
|
|
| with open(args.scfm_genes_file, "r") as handle:
|
| scfm_genes = [line.rstrip("\n") for line in handle.readlines()]
|
| scfm_gene_set = set(scfm_genes)
|
|
|
| print_sys(f"Loading data from {args.data_path}")
|
| adata = sc.read_h5ad(args.data_path)
|
|
|
| if args.min_genes_per_cell > 0:
|
| before_cells = adata.n_obs
|
| sc.pp.filter_cells(adata, min_genes=args.min_genes_per_cell)
|
| print_sys(f"Filtered cells by min_genes={args.min_genes_per_cell}: {before_cells} -> {adata.n_obs}")
|
|
|
| if args.min_cells_per_gene > 0:
|
| before_genes = adata.n_vars
|
| sc.pp.filter_genes(adata, min_cells=args.min_cells_per_gene)
|
| print_sys(f"Filtered genes by min_cells={args.min_cells_per_gene}: {before_genes} -> {adata.n_vars}")
|
|
|
| if args.n_cells == "all":
|
| n_cells = adata.n_obs
|
| else:
|
| n_cells = int(args.n_cells)
|
| if adata.n_obs > n_cells:
|
| print_sys(f"Subsampling to {n_cells} cells")
|
| sc.pp.subsample(adata, n_obs=n_cells, copy=False)
|
| n_cells = adata.n_obs
|
|
|
| batch_key = resolve_batch_key(adata, resolved_batch_key)
|
| label_keys = resolve_label_keys(adata, resolved_label_keys)
|
|
|
| print_sys(f"Processing {n_cells} cells with {adata.n_vars} genes")
|
| print_sys(f"Evaluation batch_key={batch_key}, label_keys={label_keys}")
|
|
|
| input_gene_ids = resolve_input_gene_ids(
|
| adata=adata,
|
| gene_col=resolved_gene_col,
|
| symbol_to_ensembl=symbol_to_ensembl,
|
| target_gene_set=scfm_gene_set,
|
| )
|
|
|
| input_mapping = {gene_id: index for index, gene_id in enumerate(input_gene_ids)}
|
| input_gene_num = len(input_gene_ids)
|
|
|
| intersect_genes = np.intersect1d(input_gene_ids, scfm_genes)
|
| print_sys(f"Input genes and scfm genes intersect: {len(intersect_genes)} genes")
|
|
|
| scfm_genes_pad = scfm_genes + ["PAD"]
|
| scfm_seq_len = len(scfm_genes_pad)
|
| scfm_mapping = {gene_id: index for index, gene_id in enumerate(scfm_genes_pad)}
|
|
|
| scfm_index_positions = [input_mapping.get(gene_id, input_gene_num) for gene_id in scfm_genes_pad]
|
| _ = [scfm_mapping.get(gene_id, scfm_seq_len) for gene_id in input_gene_ids]
|
|
|
| print_sys("Extracting cell embeddings...")
|
| cell_embeddings = []
|
| raw_module_embeddings = [] if args.save_raw_module_emb else None
|
| embedding_key = args.output_key
|
|
|
| with torch.no_grad():
|
| for start in tqdm(range(0, n_cells, args.batch_size)):
|
| batch_end = min(start + args.batch_size, n_cells)
|
| batch_matrix = adata.X[start:batch_end, :]
|
| if hasattr(batch_matrix, "toarray"):
|
| batch_matrix = batch_matrix.toarray()
|
| x = torch.tensor(batch_matrix, dtype=torch.float32).to(device)
|
|
|
| x_scfm = reindex_tensor_universal(x, scfm_index_positions, dim=1, filler="0", device=device)
|
| out = model(x_scfm, mask_ratio=0.0, return_encodings=True)
|
|
|
| if "module_encodings" not in out:
|
| raise ValueError(f"Expected 'module_encodings' in output, got: {list(out.keys())}")
|
|
|
| module_encodings = out["module_encodings"]
|
| ids_restore = out.get("ids_restore")
|
| if args.pool_method == "cls":
|
| if "cls_encoding" not in out:
|
| print_sys("Warning: pool_method='cls' but checkpoint has no cls_encoding, using mean pooling instead")
|
| cell_emb = module_encodings.mean(dim=1)
|
| else:
|
| cell_emb = out["cls_encoding"]
|
| elif args.pool_method == "mean":
|
| cell_emb = module_encodings.mean(dim=1)
|
| elif args.pool_method == "max":
|
| cell_emb = module_encodings.max(dim=1)[0]
|
| elif args.pool_method == "first_module":
|
| if ids_restore is None:
|
| raise ValueError("pool_method='first_module' requires ids_restore in model outputs")
|
| restore_index = ids_restore.unsqueeze(-1).expand(-1, -1, module_encodings.shape[-1])
|
| restored_module_encodings = torch.gather(module_encodings, dim=1, index=restore_index)
|
| cell_emb = restored_module_encodings[:, 0, :]
|
| else:
|
| cell_emb = out["cls_encoding"] if "cls_encoding" in out else module_encodings.mean(dim=1)
|
|
|
| if args.save_raw_module_emb:
|
| raw_module_embeddings.append(module_encodings.cpu().numpy())
|
|
|
| cell_embeddings.append(cell_emb.cpu().numpy())
|
|
|
| cell_embeddings = np.concatenate(cell_embeddings, axis=0)
|
| print_sys(f"Cell embeddings shape: {cell_embeddings.shape}")
|
|
|
| if args.save_raw_module_emb and raw_module_embeddings is not None:
|
| raw_module_embeddings = np.concatenate(raw_module_embeddings, axis=0)
|
| print_sys(f"Raw module embeddings shape: {raw_module_embeddings.shape}")
|
| raw_module_path = os.path.join(args.output_dir, "sclong_module_raw_module_emb.npy")
|
| np.save(raw_module_path, raw_module_embeddings)
|
| print_sys(f"Raw module embeddings saved to {raw_module_path}")
|
|
|
| adata.obsm[embedding_key] = cell_embeddings
|
|
|
| print_sys("Evaluating batch correction...")
|
| res = evaluate(
|
| adata=adata,
|
| batch_key=batch_key,
|
| label_keys=label_keys,
|
| embedding_key=embedding_key,
|
| metrics_mode=args.metrics_mode,
|
| res_path=os.path.join(args.output_dir, "sclong_module_batch_cell_emb_mode.csv"),
|
| )
|
|
|
| print_sys("\nEvaluation Results:")
|
| print_sys("=" * 50)
|
| print_sys(res.to_string(index=False))
|
| print_sys("=" * 50)
|
|
|
| np.save(os.path.join(args.output_dir, "sclong_module_batch_cell_emb_mode_ce.npy"), cell_embeddings)
|
| print_sys(
|
| f"Cell embeddings saved to "
|
| f"{os.path.join(args.output_dir, 'sclong_module_batch_cell_emb_mode_ce.npy')}"
|
| )
|
| print_sys("Zero-shot evaluation completed!")
|
|
|
