File size: 2,277 Bytes
94e9257 | 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 | #!/usr/bin/env python3
"""Extract L2-normalized scRep cell embeddings from an .h5ad file."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from types import SimpleNamespace
import numpy as np
RELEASE_ROOT = Path(__file__).resolve().parents[1]
if str(RELEASE_ROOT) not in sys.path:
sys.path.insert(0, str(RELEASE_ROOT))
from scRep_inference import encode_embeddings, load_examples, load_scRep_bundle
from scRep_pretrain.vocab import gene_vocab_to_map
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input_h5ad", required=True, help="Input AnnData (.h5ad) file.")
parser.add_argument("--output", default="embeddings.npy", help="Output .npy embedding matrix.")
parser.add_argument("--checkpoint", default="checkpoints/scRep_20260625_30M")
parser.add_argument("--device", default="cuda" if __import__("torch").cuda.is_available() else "cpu")
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--max_input_genes", type=int, default=2048)
parser.add_argument("--n_bins", type=int, default=50)
parser.add_argument("--max_cells", type=int, default=0, help="0 means all cells.")
parser.add_argument("--use_raw", action="store_true", help="Use adata.raw instead of adata.X.")
return parser.parse_args()
def main() -> None:
args = parse_args()
checkpoint = Path(args.checkpoint)
if not checkpoint.is_absolute():
checkpoint = RELEASE_ROOT / checkpoint
model, gene_vocab, *_ = load_scRep_bundle(checkpoint, SimpleNamespace(model_dir=str(checkpoint), asset_dir="", device=args.device))
gene_name_to_id = gene_vocab_to_map(gene_vocab)
examples, _ = load_examples(args.input_h5ad, gene_name_to_id, max_total_cells=args.max_cells, use_raw=args.use_raw)
if not examples:
raise ValueError("No input cells had genes present in the checkpoint vocabulary.")
embeddings = encode_embeddings(model, examples, args)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
np.save(output, embeddings)
print(f"Saved {embeddings.shape[0]} x {embeddings.shape[1]} embeddings to {output}")
if __name__ == "__main__":
main()
|