| """Shared RNA inputs for the RNA foundation-model extractors (43_*.py). |
| |
| load(name) -> (X, genes): X = scipy CSR float32 raw(-like) counts (n_spot, n_gene), genes = HGNC symbols (np.array str) |
| hest : HEST-1k Visium real counts (9484 spots, 14,802 common genes) -> training / test of H&E<->RNA |
| hest_v : STPath virtual RNA of the same 9484 HEST spots (calibration only) |
| orion_v : STPath virtual RNA of the 3187 ORION tiles (test only) |
| tcga_v : STPath virtual RNA of the 6144 TCGA text patches (test only) |
| dlpfc : DLPFC external Visium real counts (482 spots, 33,538 genes) -> external test only |
| STPath outputs log1p counts; virtual counts = round(expm1(pred)) so that the zero pattern means "not detected" |
| as in real Visium (models that tokenise expressed genes would otherwise see all 14,791 genes as expressed). |
| Output convention of every extractor: cache/rnafm/<model>_<set>.npy, float16 (n_spot, D), row order = the set's. |
| """ |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import scipy.sparse as sp |
|
|
| from common import CACHE |
|
|
| SETS = ["hest", "hest_v", "orion_v", "tcga_v", "dlpfc"] |
| OUT = os.path.join(CACHE, "rnafm") |
| _STP = {"hest_v": "hest_val", "orion_v": "orion", "tcga_v": "tcga_text"} |
|
|
|
|
| def load(name): |
| if name == "hest": |
| X = sp.load_npz(os.path.join(CACHE, "hest", "counts.npz")).tocsr().astype(np.float32) |
| g = pd.read_csv(os.path.join(CACHE, "hest", "genes.csv")).gene.values |
| elif name in _STP: |
| p = np.load(os.path.join(CACHE, "stpath", f"{_STP[name]}_pred.npy")).astype(np.float32) |
| X = sp.csr_matrix(np.rint(np.expm1(p).clip(0)).astype(np.float32)) |
| g = pd.read_csv(os.path.join(CACHE, "stpath", f"{_STP[name]}_genes.csv")).gene.values |
| elif name == "dlpfc": |
| X = sp.csr_matrix(np.load(os.path.join(CACHE, "external", "dlpfc_counts_full.npy")).astype(np.float32)) |
| g = pd.read_csv(os.path.join(CACHE, "external", "dlpfc_genes_full.csv")).gene.values |
| else: |
| raise KeyError(name) |
| X.eliminate_zeros() |
| return X, np.asarray(g, dtype=str) |
|
|
|
|
| def save(model, name, E): |
| os.makedirs(OUT, exist_ok=True) |
| E = np.asarray(E, dtype=np.float32) |
| assert np.isfinite(E).all(), (model, name) |
| np.save(os.path.join(OUT, f"{model}_{name}.npy"), E.astype(np.float16)) |
| print(f"[{model}] {name}: saved {E.shape}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| for s in SETS: |
| X, g = load(s) |
| nnz = np.diff(X.indptr) |
| print(f"{s:8s} {X.shape} genes/spot median {np.median(nnz):.0f} UMI/spot median {np.median(np.asarray(X.sum(1)).ravel()):.0f}") |
|
|
