File size: 2,603 Bytes
1db83f1 | 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 53 54 55 56 | """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}")
|