|
| 1 | +"""Helpers shared by ss_opm_train and ss_opm_predict.""" |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +import pandas as pd |
| 5 | +import scipy.sparse |
| 6 | + |
| 7 | +# the cell types the original ss_opm model was trained against. only used to name the |
| 8 | +# cell_ratio_* columns it expects; the ratios themselves are derived from the data when |
| 9 | +# cell type labels are available. |
| 10 | +CITE_CELL_TYPES = ["HSC", "EryP", "NeuP", "MasP", "MkP", "BP", "MoP"] |
| 11 | + |
| 12 | +# number of batch singular-vector columns the cite model expects |
| 13 | +N_BATCH_SV = 8 |
| 14 | + |
| 15 | + |
| 16 | +def to_sparse_csr(X): |
| 17 | + if scipy.sparse.issparse(X): |
| 18 | + return X.tocsr() |
| 19 | + return scipy.sparse.csr_matrix(X) |
| 20 | + |
| 21 | + |
| 22 | +def extract_day(batch, pattern=r"d(\d+)"): |
| 23 | + """Pull the day out of a batch label. |
| 24 | +
|
| 25 | + The NeurIPS 2021 batches are named `s{site}d{day}`, e.g. `s1d2`. Datasets that |
| 26 | + label their batches differently yield NaN, which the caller fills with 0 -- the |
| 27 | + model then sees a single constant day rather than failing. |
| 28 | + """ |
| 29 | + return batch.astype(str).str.extract(pattern, expand=False).astype(float) |
| 30 | + |
| 31 | + |
| 32 | +def build_metadata( |
| 33 | + adata, |
| 34 | + task_type, |
| 35 | + cell_type_col=None, |
| 36 | + group_by_batch=True, |
| 37 | + day_pattern=r"d(\d+)", |
| 38 | +): |
| 39 | + """Build the metadata frame ss_opm expects from an AnnData. |
| 40 | +
|
| 41 | + ss_opm was written against the Kaggle competition tables, which carry columns this |
| 42 | + task's API does not: `file_train_mod1.yaml` and `file_test_mod1.yaml` guarantee only |
| 43 | + `batch`. Everything else is either derived from `batch`, computed from the expression |
| 44 | + matrix, or filled with a neutral constant. |
| 45 | +
|
| 46 | + Parameters |
| 47 | + ---------- |
| 48 | + adata |
| 49 | + Input modality, with a `normalized` layer and `obs["batch"]`. |
| 50 | + task_type |
| 51 | + Either `"cite"` or `"multi"`; the cite model expects extra columns. |
| 52 | + cell_type_col |
| 53 | + Column in `adata.obs` holding cell type labels. When given, it drives both |
| 54 | + `cell_type` and the `cell_ratio_*` columns. When None -- the case for every |
| 55 | + dataset this task currently ships -- cell types are `"hidden"` and the ratios |
| 56 | + are uniform. |
| 57 | + group_by_batch |
| 58 | + Assign one group per batch. Set False to put every cell in group 0, which is |
| 59 | + what the predict path wants, since targets are absent and the group IDs are |
| 60 | + only used to look up target statistics. |
| 61 | + day_pattern |
| 62 | + Regex whose first capture group is the day within a batch label. |
| 63 | + """ |
| 64 | + obs = pd.DataFrame(index=adata.obs_names) |
| 65 | + |
| 66 | + obs["batch"] = adata.obs["batch"].values |
| 67 | + obs["day"] = extract_day(adata.obs["batch"], day_pattern).fillna(0).values |
| 68 | + |
| 69 | + # per-cell statistics from the normalized expression layer |
| 70 | + X = adata.layers["normalized"] |
| 71 | + X_dense = X.toarray() if scipy.sparse.issparse(X) else np.asarray(X, dtype=float) |
| 72 | + |
| 73 | + obs["nonzero_ratio"] = (X_dense != 0).mean(axis=1) |
| 74 | + obs["nonzero_q25"] = np.percentile(X_dense, 25, axis=1) |
| 75 | + obs["nonzero_q50"] = np.percentile(X_dense, 50, axis=1) |
| 76 | + obs["nonzero_q75"] = np.percentile(X_dense, 75, axis=1) |
| 77 | + obs["mean"] = X_dense.mean(axis=1) |
| 78 | + obs["std"] = X_dense.std(axis=1) |
| 79 | + |
| 80 | + if group_by_batch: |
| 81 | + batches = adata.obs["batch"].unique().tolist() |
| 82 | + obs["group"] = adata.obs["batch"].map({b: i for i, b in enumerate(batches)}).astype(int).values |
| 83 | + else: |
| 84 | + obs["group"] = 0 |
| 85 | + |
| 86 | + # cell type labels, when the caller can supply them |
| 87 | + if cell_type_col is not None and cell_type_col in adata.obs: |
| 88 | + obs["cell_type"] = adata.obs[cell_type_col].astype(str).values |
| 89 | + else: |
| 90 | + obs["cell_type"] = "hidden" |
| 91 | + |
| 92 | + # donor and technology are not in this task's file format; gender_id defaults to 0 |
| 93 | + obs["donor"] = 0 |
| 94 | + obs["technology"] = "unknown" |
| 95 | + |
| 96 | + if task_type == "cite": |
| 97 | + ratios = obs["cell_type"].value_counts(normalize=True) |
| 98 | + for cell_type in CITE_CELL_TYPES: |
| 99 | + obs[f"cell_ratio_{cell_type}"] = ratios.get(cell_type, 1.0 / len(CITE_CELL_TYPES)) |
| 100 | + |
| 101 | + batch_counts = adata.obs["batch"].value_counts() |
| 102 | + obs["cell_count"] = adata.obs["batch"].map(batch_counts).astype(float).values |
| 103 | + |
| 104 | + # the originals are singular vectors of the full Kaggle batch matrix, which we |
| 105 | + # cannot reconstruct from a single dataset |
| 106 | + for i in range(N_BATCH_SV): |
| 107 | + obs[f"batch_sv{i}"] = 0.0 |
| 108 | + |
| 109 | + return obs |
0 commit comments