| """AnnData to PyTorch DataLoader conversion for DeepPTR.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| if TYPE_CHECKING: |
| from anndata import AnnData |
|
|
|
|
| def setup_dataloaders( |
| adata: AnnData, |
| batch_size: int = 256, |
| val_frac: float = 0.1, |
| stratify_key: str | None = None, |
| seed: int = 0, |
| num_workers: int = 0, |
| ) -> tuple[DataLoader, DataLoader, np.ndarray, np.ndarray]: |
| """Build train and validation DataLoaders from an AnnData object. |
| |
| Extracts raw integer counts from ``adata.layers['spliced']`` and |
| ``adata.layers['unspliced']``, together with per-cell library sizes. |
| |
| Parameters |
| ---------- |
| adata |
| Annotated data matrix. Must contain ``layers['spliced']`` and |
| ``layers['unspliced']``. |
| batch_size |
| Mini-batch size. |
| val_frac |
| Fraction of cells held out for validation. |
| stratify_key |
| Optional obs column for stratified splitting. |
| seed |
| Random seed for reproducibility. |
| num_workers |
| DataLoader workers. |
| |
| Returns |
| ------- |
| train_dl, val_dl, train_idx, val_idx |
| """ |
| from scipy.sparse import issparse |
|
|
| from ._utils import get_library_sizes |
|
|
| def _dense(mat): |
| if issparse(mat): |
| return np.asarray(mat.todense()) |
| return np.asarray(mat) |
|
|
| s = _dense(adata.layers["spliced"]).astype(np.float32) |
| u = _dense(adata.layers["unspliced"]).astype(np.float32) |
| l_u, l_s = get_library_sizes(adata) |
|
|
| n = adata.n_obs |
| indices = np.arange(n) |
|
|
| if stratify_key is not None and stratify_key in adata.obs.columns: |
| from sklearn.model_selection import StratifiedShuffleSplit |
|
|
| labels = adata.obs[stratify_key].values |
| splitter = StratifiedShuffleSplit( |
| n_splits=1, test_size=val_frac, random_state=seed |
| ) |
| train_idx, val_idx = next(splitter.split(indices, labels)) |
| else: |
| rng = np.random.RandomState(seed) |
| perm = rng.permutation(n) |
| n_val = max(1, int(n * val_frac)) |
| val_idx = perm[:n_val] |
| train_idx = perm[n_val:] |
|
|
| def _make_loader(idx: np.ndarray, shuffle: bool) -> DataLoader: |
| ds = TensorDataset( |
| torch.from_numpy(s[idx]), |
| torch.from_numpy(u[idx]), |
| torch.from_numpy(l_s[idx]), |
| torch.from_numpy(l_u[idx]), |
| ) |
| return DataLoader( |
| ds, |
| batch_size=batch_size, |
| shuffle=shuffle, |
| num_workers=num_workers, |
| pin_memory=False, |
| drop_last=False, |
| ) |
|
|
| train_dl = _make_loader(train_idx, shuffle=True) |
| val_dl = _make_loader(val_idx, shuffle=False) |
| return train_dl, val_dl, train_idx, val_idx |
|
|
