scdataloader 1.9.2__py3-none-any.whl → 2.0.2__py3-none-any.whl

scdataloader/mapped.py

@@ -7,12 +7,14 @@
 
 from __future__ import annotations
 
+import os
 from collections import Counter
 from functools import reduce
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, List, Literal
 
 import numpy as np
 import pandas as pd
+import torch
 from lamindb.core.storage._anndata_accessor import (
     ArrayType,
     ArrayTypes,
@@ -24,10 +26,13 @@ from lamindb.core.storage._anndata_accessor import (
     registry,
 )
 from lamindb_setup.core.upath import UPath
+from tqdm import tqdm
 
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
+from pandas.api.types import union_categoricals
+
 
 class _Connect:
     def __init__(self, storage):
@@ -106,24 +111,28 @@ class MappedCollection:
         meta_assays: Assays that are already defined as metacells.
         metacell_mode: frequency at which to sample a metacell (an average of k-nearest neighbors).
         get_knn_cells: Whether to also dataload the k-nearest neighbors of each queried cells.
+        store_location: Path to a directory where klass_indices can be cached, or full path to the cache file.
+        force_recompute_indices: If True, recompute indices even if a cache file exists.
     """
 
     def __init__(
         self,
-        path_list: list[UPathStr],
-        layers_keys: str | list[str] | None = None,
-        obs_keys: str | list[str] | None = None,
-        obsm_keys: str | list[str] | None = None,
+        path_list: List[UPathStr],
+        layers_keys: str | List[str] | None = None,
+        obs_keys: str | List[str] | None = None,
+        obsm_keys: str | List[str] | None = None,
         obs_filter: dict[str, str | tuple[str, ...]] | None = None,
         join: Literal["inner", "outer"] | None = "inner",
-        encode_labels: bool | list[str] = True,
+        encode_labels: bool | List[str] = True,
         unknown_label: str | dict[str, str] | None = None,
         cache_categories: bool = True,
         parallel: bool = False,
         dtype: str | None = None,
         metacell_mode: float = 0.0,
         get_knn_cells: bool = False,
-        meta_assays: list[str] = ["EFO:0022857", "EFO:0010961"],
+        meta_assays: List[str] = ["EFO:0022857", "EFO:0010961"],
+        store_location: str | None = None,
+        force_recompute_indices: bool = False,
     ):
         if join not in {None, "inner", "outer"}:  # pragma: nocover
             raise ValueError(
@@ -181,14 +190,28 @@ class MappedCollection:
         self._cache_cats: dict = {}
         if self.obs_keys is not None:
             if cache_categories:
-                self._cache_categories(self.obs_keys)
+                if store_location is not None:
+                    os.makedirs(store_location, exist_ok=True)
+                    self.store_location = os.path.join(store_location, "categories")
+                    if (
+                        not os.path.exists(self.store_location)
+                        or force_recompute_indices
+                    ):
+                        self._cache_categories(self.obs_keys)
+                        torch.save(self._cache_cats, self.store_location)
+                    else:
+                        self._cache_cats = torch.load(
+                            self.store_location, weights_only=False
+                        )
+                        print(f"Loaded categories from {self.store_location}")
             self.encoders: dict = {}
             if self.encode_labels:
                 self._make_encoders(self.encode_labels)  # type: ignore
-
         self.n_obs_list = []
         self.indices_list = []
-        for i, storage in enumerate(self.storages):
+        for i, storage in tqdm(
+            enumerate(self.storages), desc="Checking datasets", total=len(self.storages)
+        ):
             with _Connect(storage) as store:
                 X = store["X"]
                 store_path = self.path_list[i]
@@ -263,13 +286,10 @@ class MappedCollection:
         self._cache_cats = {}
         for label in obs_keys:
             self._cache_cats[label] = []
-            for storage in self.storages:
+            for storage in tqdm(self.storages, f"caching categories, {label}"):
                 with _Connect(storage) as store:
                     cats = self._get_categories(store, label)
-                    if cats is not None:
-                        cats = (
-                            _decode(cats) if isinstance(cats[0], bytes) else cats[...]
-                        )
+                    cats = _decode(cats) if isinstance(cats[0], bytes) else cats[...]
                     self._cache_cats[label].append(cats)
 
     def _make_encoders(self, encode_labels: list):
@@ -330,7 +350,7 @@ class MappedCollection:
             vrs_sort_status = (vrs.is_monotonic_decreasing for vrs in self.var_list)
         return all(vrs_sort_status)
 
-    def check_vars_non_aligned(self, vars: pd.Index | list) -> list[int]:
+    def check_vars_non_aligned(self, vars: pd.Index | List) -> List[int]:
         """Returns indices of objects with non-aligned variables.
 
         Args:
@@ -362,7 +382,7 @@ class MappedCollection:
         return (self.n_obs, self.n_vars)
 
     @property
-    def original_shapes(self) -> list[tuple[int, int]]:
+    def original_shapes(self) -> List[tuple[int, int]]:
         """Shapes of the underlying AnnData objects (with `obs_filter` applied)."""
         if self.n_vars_list is None:
             n_vars_list = [None] * len(self.n_obs_list)
@@ -403,10 +423,36 @@ class MappedCollection:
                         cats = None
                     label_idx = self._get_obs_idx(store, obs_idx, label, cats)
                     if label in self.encoders:
-                        label_idx = self.encoders[label][label_idx]
-                    out[label] = label_idx
-
-            if self.metacell_mode > 0:
+                        try:
+                            label_idx = self.encoders[label][label_idx]
+                        except:
+                            print(self.storages[storage_idx])
+                            print(label, label_idx)
+                            print(idx)
+                            print(cats)
+                            raise
+                    try:
+                        out[label] = label_idx
+                    except:
+                        print(self.storages[storage_idx])
+                        print(label, label_idx)
+                        print(out)
+                        raise
+
+            if self.get_knn_cells:
+                distances = self._get_data_idx(store["obsp"]["distances"], obs_idx)
+                nn_idx = np.argsort(-1 / (distances - 1e-6))[:6]
+                out["knn_cells"] = np.array(
+                    [
+                        self._get_data_idx(
+                            lazy_data, i, self.join_vars, var_idxs_join, self.n_vars
+                        )
+                        for i in nn_idx
+                    ],
+                    dtype=int,
+                )
+                out["knn_cells_info"] = distances[nn_idx]
+            elif self.metacell_mode > 0:
                 if (
                     len(self.meta_assays) > 0
                     and "assay_ontology_term_id" in self.obs_keys
@@ -423,19 +469,6 @@ class MappedCollection:
                         out[layers_key] += self._get_data_idx(
                             lazy_data, i, self.join_vars, var_idxs_join, self.n_vars
                         )
-            elif self.get_knn_cells:
-                distances = self._get_data_idx(store["obsp"]["distances"], obs_idx)
-                nn_idx = np.argsort(-1 / (distances - 1e-6))[:6]
-                out["knn_cells"] = np.array(
-                    [
-                        self._get_data_idx(
-                            lazy_data, i, self.join_vars, var_idxs_join, self.n_vars
-                        )
-                        for i in nn_idx
-                    ],
-                    dtype=int,
-                )
-                out["distances"] = distances[nn_idx]
 
         return out
 
@@ -510,7 +543,7 @@ class MappedCollection:
 
     def get_label_weights(
         self,
-        obs_keys: str | list[str],
+        obs_keys: str | List[str],
         scaler: float | None = None,
         return_categories: bool = False,
     ):
@@ -555,21 +588,41 @@ class MappedCollection:
             weights = (MAX / scaler) / ((1 + counts - MIN) + MAX / scaler)
         return weights
 
-    def get_merged_labels(self, label_key: str):
+    def get_merged_labels(self, label_key: str, is_cat: bool = True):
         """Get merged labels for `label_key` from all `.obs`."""
         labels_merge = []
-        for i, storage in enumerate(self.storages):
+        for i, storage in tqdm(
+            enumerate(self.storages), label_key, total=len(self.storages)
+        ):
             with _Connect(storage) as store:
-                labels = self._get_labels(store, label_key, storage_idx=i)
+                labels = self._get_labels(
+                    store, label_key, storage_idx=i, is_cat=is_cat
+                )
                 if self.filtered:
                     labels = labels[self.indices_list[i]]
                 labels_merge.append(labels)
-        return np.hstack(labels_merge)
+        if is_cat:
+            try:
+                return union_categoricals(labels_merge)
+            except TypeError:
+                typ = type(int)
+                for i in range(len(labels_merge)):
+                    if typ != type(labels_merge[i][0]):
+                        self.storages[i]
+                    typ = type(labels_merge[i][0])
+                return []
+        else:
+            print("concatenating labels")
+            return np.concatenate(labels_merge)
 
     def get_merged_categories(self, label_key: str):
         """Get merged categories for `label_key` from all `.obs`."""
         cats_merge = set()
-        for i, storage in enumerate(self.storages):
+        for i, storage in tqdm(
+            enumerate(self.storages),
+            total=len(self.storages),
+            desc="merging all " + label_key + " categories",
+        ):
             with _Connect(storage) as store:
                 if label_key in self._cache_cats:
                     cats = self._cache_cats[label_key][i]
@@ -609,8 +662,8 @@ class MappedCollection:
             else:
                 if "categories" in labels.attrs:
                     return labels.attrs["categories"]
-                else:
-                    return None
+                elif labels.dtype == "bool":
+                    return np.array(["True", "False"])
         return None
 
     def _get_codes(self, storage: StorageType, label_key: str):
@@ -626,11 +679,17 @@ class MappedCollection:
                 return label["codes"][...]
 
     def _get_labels(
-        self, storage: StorageType, label_key: str, storage_idx: int | None = None
+        self,
+        storage: StorageType,
+        label_key: str,
+        storage_idx: int | None = None,
+        is_cat: bool = True,
     ):
         """Get labels."""
         codes = self._get_codes(storage, label_key)
         labels = _decode(codes) if isinstance(codes[0], bytes) else codes
+        if labels.dtype == bool:
+            labels = labels.astype(int)
         if storage_idx is not None and label_key in self._cache_cats:
             cats = self._cache_cats[label_key][storage_idx]
         else:
@@ -638,6 +697,8 @@ class MappedCollection:
         if cats is not None:
             cats = _decode(cats) if isinstance(cats[0], bytes) else cats
             labels = cats[labels]
+        if is_cat:
+            labels = pd.Categorical(labels.astype(str))
         return labels
 
     def close(self):