scdataloader 1.9.2__py3-none-any.whl → 2.0.0__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
 
 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,6 +111,8 @@ 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__(
@@ -124,6 +131,8 @@ class MappedCollection:
         metacell_mode: float = 0.0,
         get_knn_cells: bool = False,
         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,26 @@ 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)
+                        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 +284,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):
@@ -403,8 +421,21 @@ 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
+                        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.metacell_mode > 0:
                 if (
@@ -555,21 +586,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 +660,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 +677,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 +695,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):

scdataloader/preprocess.py

@@ -1,4 +1,5 @@
 import gc
+import time
 from typing import Callable, Optional, Union
 from uuid import uuid4
 
@@ -8,6 +9,7 @@ import numpy as np
 import pandas as pd
 import scanpy as sc
 from anndata import AnnData, read_h5ad
+from django.db.utils import OperationalError
 from scipy.sparse import csr_matrix
 from upath import UPath
 
@@ -61,6 +63,7 @@ class Preprocessor:
         organisms: list[str] = ["NCBITaxon:9606", "NCBITaxon:10090"],
         use_raw: bool = True,
         keepdata: bool = False,
+        drop_non_primary: bool = False,
     ) -> None:
         """
         Initializes the preprocessor and configures the workflow steps.
@@ -108,6 +111,8 @@ class Preprocessor:
                 Defaults to False.
             keepdata (bool, optional): Determines whether to keep the data in the AnnData object.
                 Defaults to False.
+            drop_non_primary (bool, optional): Determines whether to drop non-primary cells.
+                Defaults to False.
         """
         self.filter_gene_by_counts = filter_gene_by_counts
         self.filter_cell_by_counts = filter_cell_by_counts
@@ -123,6 +128,7 @@ class Preprocessor:
         self.min_valid_genes_id = min_valid_genes_id
         self.min_nnz_genes = min_nnz_genes
         self.maxdropamount = maxdropamount
+        self.drop_non_primary = drop_non_primary
         self.madoutlier = madoutlier
         self.n_hvg_for_postp = n_hvg_for_postp
         self.pct_mt_outlier = pct_mt_outlier
@@ -142,10 +148,6 @@ class Preprocessor:
             raise ValueError(
                 "organism_ontology_term_id not found in adata.obs, you need to add an ontology term id for the organism of your anndata"
             )
-        if not adata[0].var.index.str.contains("ENS").any() and not self.is_symbol:
-            raise ValueError(
-                "gene names in the `var.index` field of your anndata should map to the ensembl_gene nomenclature else set `is_symbol` to True if using hugo symbols"
-            )
         if adata.obs["organism_ontology_term_id"].iloc[0] not in self.organisms:
             raise ValueError(
                 "we cannot work with this organism",
@@ -161,8 +163,8 @@ class Preprocessor:
                 if np.abs(adata[:50_000].X.astype(int) - adata[:50_000].X).sum():
                     print("X was not raw counts, using 'counts' layer")
                     adata.X = adata.layers["counts"].copy()
-            print("Dropping layers: ", adata.layers.keys())
             if not self.keepdata:
+                print("Dropping layers: ", adata.layers.keys())
                 del adata.layers
         if len(adata.varm.keys()) > 0 and not self.keepdata:
             del adata.varm
@@ -170,6 +172,8 @@ class Preprocessor:
             del adata.obsm
         if len(adata.obsp.keys()) > 0 and not self.keepdata:
             del adata.obsp
+        if len(adata.varp.keys()) > 0 and not self.keepdata:
+            del adata.varp
         # check that it is a count
 
         print("checking raw counts")
@@ -188,7 +192,7 @@ class Preprocessor:
             # if not available count drop
         prevsize = adata.shape[0]
         # dropping non primary
-        if "is_primary_data" in adata.obs.columns:
+        if "is_primary_data" in adata.obs.columns and self.drop_non_primary:
             adata = adata[adata.obs.is_primary_data]
         if adata.shape[0] < self.min_dataset_size:
             raise Exception("Dataset dropped due to too many secondary cells")
@@ -213,13 +217,10 @@ class Preprocessor:
                 min_genes=self.min_nnz_genes,
             )
         # if lost > 50% of the dataset, drop dataset
-        # load the genes
-        genesdf = data_utils.load_genes(adata.obs.organism_ontology_term_id.iloc[0])
-
-        if prevsize / adata.shape[0] > self.maxdropamount:
+        if prevsize / (adata.shape[0] + 1) > self.maxdropamount:
             raise Exception(
                 "Dataset dropped due to low expressed genes and unexpressed cells: factor of "
-                + str(prevsize / adata.shape[0])
+                + str(prevsize / (adata.shape[0] + 1))
             )
         if adata.shape[0] < self.min_dataset_size:
             raise Exception(
@@ -232,60 +233,39 @@ class Preprocessor:
             )
         )
 
-        # Check if we have a mix of gene names and ensembl IDs
-        has_ens = adata.var.index.str.match(r"ENS.*\d{6,}$").any()
-        all_ens = adata.var.index.str.match(r"ENS.*\d{6,}$").all()
-
-        if not has_ens:
-            print("No ENS genes found, assuming gene symbols...")
-        elif not all_ens:
-            print("Mix of ENS and gene symbols found, converting all to ENS IDs...")
-
+        # load the genes
+        genesdf = data_utils.load_genes(adata.obs.organism_ontology_term_id.iloc[0])
         genesdf["ensembl_gene_id"] = genesdf.index
 
         # For genes that are already ENS IDs, use them directly
-        ens_mask = adata.var.index.str.match(r"ENS.*\d{6,}$")
-        symbol_mask = ~ens_mask
-
+        prev_size = adata.shape[1]
         # Handle symbol genes
-        if symbol_mask.any():
-            symbol_var = adata.var[symbol_mask].merge(
+        if self.is_symbol:
+            new_var = adata.var.merge(
                 genesdf.drop_duplicates("symbol").set_index("symbol", drop=False),
                 left_index=True,
                 right_index=True,
                 how="inner",
             )
-
-        # Handle ENS genes
-        if ens_mask.any():
-            ens_var = adata.var[ens_mask].merge(
+            new_var["symbol"] = new_var.index
+            adata = adata[:, new_var.index]
+            new_var.index = new_var["ensembl_gene_id"]
+        else:
+            new_var = adata.var.merge(
                 genesdf, left_index=True, right_index=True, how="inner"
             )
+            adata = adata[:, new_var.index]
+        print(f"Removed {prev_size - adata.shape[1]} genes not known to the ontology")
+        prev_size = adata.shape[1]
 
-        # Combine and sort
-        if symbol_mask.any() and ens_mask.any():
-            var = pd.concat([symbol_var, ens_var])
-        elif symbol_mask.any():
-            var = symbol_var
-        else:
-            var = ens_var
-
-        adata = adata[:, var.index]
-        #        var = var.sort_values(by="ensembl_gene_id").set_index("ensembl_gene_id")
-        # Update adata with combined genes
-        if "ensembl_gene_id" in var.columns:
-            adata.var = var.set_index("ensembl_gene_id")
-        else:
-            adata.var = var
+        adata.var = new_var
         # Drop duplicate genes, keeping first occurrence
         adata = adata[:, ~adata.var.index.duplicated(keep="first")]
+        print(f"Removed {prev_size - adata.shape[1]} duplicate genes")
 
-        intersect_genes = set(adata.var.index).intersection(set(genesdf.index))
-        print(f"Removed {len(adata.var.index) - len(intersect_genes)} genes.")
-        if len(intersect_genes) < self.min_valid_genes_id:
+        if adata.shape[1] < self.min_valid_genes_id:
             raise Exception("Dataset dropped due to too many genes not mapping to it")
-        adata = adata[:, list(intersect_genes)]
-        # marking unseen genes
+
         unseen = set(genesdf.index) - set(adata.var.index)
         # adding them to adata
         emptyda = ad.AnnData(
@@ -293,6 +273,9 @@ class Preprocessor:
             var=pd.DataFrame(index=list(unseen)),
             obs=pd.DataFrame(index=adata.obs.index),
         )
+        print(
+            f"Added {len(unseen)} genes in the ontology but not present in the dataset"
+        )
         adata = ad.concat([adata, emptyda], axis=1, join="outer", merge="only")
         # do a validation function
         adata.uns["unseen_genes"] = list(unseen)
@@ -330,7 +313,7 @@ class Preprocessor:
         # QC
 
         adata.var[genesdf.columns] = genesdf.loc[adata.var.index]
-        print("startin QC")
+        print("starting QC")
         sc.pp.calculate_qc_metrics(
             adata, qc_vars=["mt", "ribo", "hb"], inplace=True, percent_top=[20]
         )
@@ -348,7 +331,7 @@ class Preprocessor:
         )
         total_outliers = (adata.obs["outlier"] | adata.obs["mt_outlier"]).sum()
         total_cells = adata.shape[0]
-        percentage_outliers = (total_outliers / total_cells) * 100
+        percentage_outliers = (total_outliers / (total_cells + 1)) * 100
         print(
             f"Seeing {total_outliers} outliers ({percentage_outliers:.2f}% of total dataset):"
         )
@@ -395,7 +378,7 @@ class Preprocessor:
                         subset=False,
                         layer="norm",
                     )
-
+            print("starting PCA")
             adata.obsm["X_pca"] = sc.pp.pca(
                 adata.layers["norm"][:, adata.var.highly_variable]
                 if "highly_variable" in adata.var.columns
@@ -464,13 +447,13 @@ class LaminPreprocessor(Preprocessor):
         *args,
         cache: bool = True,
         keep_files: bool = True,
-        force_preloaded: bool = False,
+        force_lamin_cache: bool = False,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.cache = cache
         self.keep_files = keep_files
-        self.force_preloaded = force_preloaded
+        self.force_lamin_cache = force_lamin_cache
 
     def __call__(
         self,
@@ -505,10 +488,13 @@ class LaminPreprocessor(Preprocessor):
                     print(f"{file.stem_uid} is already processed... not preprocessing")
                     continue
                 print(file)
+                if self.force_lamin_cache:
+                    path = cache_path(file)
+                    backed = read_h5ad(path, backed="r")
+                else:
+                    # file.cache()
+                    backed = file.open()
 
-                _ = cache_path(file) if self.force_preloaded else file.cache()
-                backed = file.open()
-                # backed = read_h5ad(path, backed="r")
                 if "is_primary_data" in backed.obs.columns:
                     if backed.obs.is_primary_data.sum() == 0:
                         print(f"{file.key} only contains non primary cells.. dropping")
@@ -556,37 +542,52 @@ class LaminPreprocessor(Preprocessor):
                                 block,
                                 dataset_id=file.stem_uid + "_p" + str(j),
                             )
-                            myfile = ln.Artifact.from_anndata(
-                                block,
-                                description=description
-                                + " n"
-                                + str(i)
-                                + " p"
-                                + str(j)
-                                + " ( revises file "
-                                + str(file.stem_uid)
-                                + " )",
-                                version=version,
-                            )
-                            myfile.save()
-
+                            saved = False
+                            while not saved:
+                                try:
+                                    myfile = ln.Artifact.from_anndata(
+                                        block,
+                                        description=description
+                                        + " n"
+                                        + str(i)
+                                        + " p"
+                                        + str(j)
+                                        + " ( revises file "
+                                        + str(file.stem_uid)
+                                        + " )",
+                                        version=version,
+                                    )
+                                    myfile.save()
+                                    saved = True
+                                except OperationalError:
+                                    print(
+                                        "Database locked, waiting 30 seconds and retrying..."
+                                    )
+                                    time.sleep(10)
                             if self.keep_files:
                                 files.append(myfile)
                                 del block
                             else:
                                 del myfile
                                 del block
-                            gc.collect()
-
                     else:
                         adata = super().__call__(adata, dataset_id=file.stem_uid)
-                        myfile = ln.Artifact.from_anndata(
-                            adata,
-                            revises=file,
-                            description=description + " p" + str(i),
-                            version=version,
-                        )
-                        myfile.save()
+                        saved = False
+                        while not saved:
+                            try:
+                                myfile = ln.Artifact.from_anndata(
+                                    adata,
+                                    revises=file,
+                                    description=description + " p" + str(i),
+                                    version=version,
+                                )
+                                myfile.save()
+                                saved = True
+                            except OperationalError:
+                                print(
+                                    "Database locked, waiting 10 seconds and retrying..."
+                                )
+                                time.sleep(10)
                         if self.keep_files:
                             files.append(myfile)
                             del adata
@@ -606,7 +607,7 @@ class LaminPreprocessor(Preprocessor):
                         continue
                     else:
                         raise e
-
+                gc.collect()
                 # issues with KLlggfw6I6lvmbqiZm46
             if self.keep_files:
                 # Reconstruct collection using keys
@@ -716,7 +717,7 @@ def additional_preprocess(adata):
             }
         }
     )  # multi ethnic will have to get renamed
-    adata.obs["cell_culture"] = False
+    adata.obs["cell_culture"] = "False"
     # if cell_type contains the word "(cell culture)" then it is a cell culture and we mark it as so and remove this from the cell type
     loc = adata.obs["cell_type_ontology_term_id"].str.contains(
         "(cell culture)", regex=False
@@ -725,7 +726,7 @@ def additional_preprocess(adata):
         adata.obs["cell_type_ontology_term_id"] = adata.obs[
             "cell_type_ontology_term_id"
         ].astype(str)
-        adata.obs.loc[loc, "cell_culture"] = True
+        adata.obs.loc[loc, "cell_culture"] = "True"
         adata.obs.loc[loc, "cell_type_ontology_term_id"] = adata.obs.loc[
             loc, "cell_type_ontology_term_id"
         ].str.replace(" (cell culture)", "")
@@ -734,7 +735,7 @@ def additional_preprocess(adata):
         "(cell culture)", regex=False
     )
     if loc.sum() > 0:
-        adata.obs.loc[loc, "cell_culture"] = True
+        adata.obs.loc[loc, "cell_culture"] = "True"
         adata.obs["tissue_ontology_term_id"] = adata.obs[
             "tissue_ontology_term_id"
         ].astype(str)
@@ -744,7 +745,7 @@ def additional_preprocess(adata):
 
     loc = adata.obs["tissue_ontology_term_id"].str.contains("(organoid)", regex=False)
     if loc.sum() > 0:
-        adata.obs.loc[loc, "cell_culture"] = True
+        adata.obs.loc[loc, "cell_culture"] = "True"
         adata.obs["tissue_ontology_term_id"] = adata.obs[
             "tissue_ontology_term_id"
         ].astype(str)
@@ -773,6 +774,7 @@ def additional_postprocess(adata):
     #    sc.external.pp.harmony_integrate(adata, key="batches")
     #    sc.pp.neighbors(adata, use_rep="X_pca_harmony")
     # else:
+    print("starting post processing")
     sc.pp.neighbors(adata, use_rep="X_pca")
     sc.tl.leiden(adata, key_added="leiden_2", resolution=2.0)
     sc.tl.leiden(adata, key_added="leiden_1", resolution=1.0)
@@ -791,8 +793,12 @@ def additional_postprocess(adata):
     MAXSIM = 0.94
     from collections import Counter
 
+    import bionty as bt
+
     from .config import MAIN_HUMAN_MOUSE_DEV_STAGE_MAP
 
+    remap_stages = {u: k for k, v in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP.items() for u in v}
+
     adata.obs[NEWOBS] = (
         adata.obs[COL].astype(str) + "_" + adata.obs["leiden_1"].astype(str)
     )
@@ -860,18 +866,17 @@ def additional_postprocess(adata):
                 num += 1
     adata.obs[NEWOBS] = adata.obs[NEWOBS].map(merge_mapping).fillna(adata.obs[NEWOBS])
 
-    import bionty as bt
-
     stages = adata.obs["development_stage_ontology_term_id"].unique()
     if adata.obs.organism_ontology_term_id.unique() == ["NCBITaxon:9606"]:
         relabel = {i: i for i in stages}
         for stage in stages:
+            if stage in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP.keys():
+                continue
             stage_obj = bt.DevelopmentalStage.filter(ontology_id=stage).first()
             parents = set([i.ontology_id for i in stage_obj.parents.filter()])
             parents = parents - set(
                 [
                     "HsapDv:0010000",
-                    "HsapDv:0000204",
                     "HsapDv:0000227",
                 ]
             )
@@ -879,9 +884,14 @@ def additional_postprocess(adata):
                 for p in parents:
                     if p in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP:
                         relabel[stage] = p
-        adata.obs["simplified_dev_stage"] = adata.obs[
-            "development_stage_ontology_term_id"
-        ].map(relabel)
+        adata.obs["age_group"] = adata.obs["development_stage_ontology_term_id"].map(
+            relabel
+        )
+        for stage in adata.obs["age_group"].unique():
+            if stage in remap_stages.keys():
+                adata.obs["age_group"] = adata.obs["age_group"].map(
+                    lambda x: remap_stages[x] if x == stage else x
+                )
     elif adata.obs.organism_ontology_term_id.unique() == ["NCBITaxon:10090"]:
         rename_mapping = {
             k: v for v, j in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP.items() for k in j
@@ -890,11 +900,12 @@ def additional_postprocess(adata):
         for stage in stages:
             if stage in rename_mapping:
                 relabel[stage] = rename_mapping[stage]
-        adata.obs["simplified_dev_stage"] = adata.obs[
-            "development_stage_ontology_term_id"
-        ].map(relabel)
+        adata.obs["age_group"] = adata.obs["development_stage_ontology_term_id"].map(
+            relabel
+        )
     else:
-        raise ValueError("organism not supported")
+        # raise ValueError("organism not supported")
+        print("organism not supported for age labels")
     # palantir.utils.run_diffusion_maps(adata, n_components=20)
     # palantir.utils.determine_multiscale_space(adata)
     # terminal_states = palantir.utils.find_terminal_states(