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

scdataloader/preprocess.py

@@ -1,5 +1,6 @@
 import gc
-from typing import Callable, Optional, Union
+import time
+from typing import Callable, List, Optional, Union
 from uuid import uuid4
 
 import anndata as ad
@@ -8,13 +9,14 @@ 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
 
 from scdataloader import utils as data_utils
 
 FULL_LENGTH_ASSAYS = [
-    "EFO: 0700016",
+    "EFO:0700016",
     "EFO:0008930",
     "EFO:0008931",
 ]
@@ -47,20 +49,21 @@ class Preprocessor:
         maxdropamount: int = 50,
         madoutlier: int = 5,
         pct_mt_outlier: int = 8,
-        batch_keys: list[str] = [
+        batch_keys: List[str] = [
             "assay_ontology_term_id",
             "self_reported_ethnicity_ontology_term_id",
             "sex_ontology_term_id",
             "donor_id",
             "suspension_type",
         ],
-        skip_validate: bool = False,
+        skip_validate: bool = True,
         additional_preprocess: Optional[Callable[[AnnData], AnnData]] = None,
         additional_postprocess: Optional[Callable[[AnnData], AnnData]] = None,
         do_postp: bool = True,
-        organisms: list[str] = ["NCBITaxon:9606", "NCBITaxon:10090"],
+        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
@@ -139,13 +145,14 @@ class Preprocessor:
         if self.additional_preprocess is not None:
             adata = self.additional_preprocess(adata)
         if "organism_ontology_term_id" not in adata[0].obs.columns:
-            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 "organism_ontology_term_id" in adata.uns:
+                adata.obs["organism_ontology_term_id"] = adata.uns[
+                    "organism_ontology_term_id"
+                ]
+            else:
+                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 adata.obs["organism_ontology_term_id"].iloc[0] not in self.organisms:
             raise ValueError(
                 "we cannot work with this organism",
@@ -161,8 +168,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 +177,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 +197,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 +222,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 +238,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 +278,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 +318,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 +336,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):"
         )
@@ -374,6 +362,8 @@ class Preprocessor:
             adata.obs["batches"] = adata.obs[batches].apply(
                 lambda x: ",".join(x.dropna().astype(str)), axis=1
             )
+            if "highly_variable" in adata.var.columns:
+                adata.var = adata.var.drop(columns=["highly_variable"])
             if self.n_hvg_for_postp:
                 try:
                     sc.pp.highly_variable_genes(
@@ -395,12 +385,15 @@ 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
-                else adata.layers["norm"],
-                n_comps=200 if adata.shape[0] > 200 else adata.shape[0] - 2,
+                (
+                    adata.layers["norm"][:, adata.var["highly_variable"]]
+                    if "highly_variable" in adata.var.columns
+                    else adata.layers["norm"]
+                ),
+                n_comps=50 if adata.shape[0] > 1000 else adata.shape[0] // 20,
+                chunked=adata.shape[0] > 100_000,
             )
 
             # additional
@@ -464,13 +457,15 @@ class LaminPreprocessor(Preprocessor):
         *args,
         cache: bool = True,
         keep_files: bool = True,
-        force_preloaded: bool = False,
+        force_lamin_cache: bool = False,
+        assays_to_drop: List[str] = ["EFO:0008939"],
         **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
+        self.assays_to_drop = assays_to_drop
 
     def __call__(
         self,
@@ -505,19 +500,25 @@ 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")
                         # Save the stem_uid to a file to avoid loading it again
-                    with open("nonprimary.txt", "a") as f:
-                        f.write(f"{file.stem_uid}\n")
-                    continue
+                        with open("nonprimary.txt", "a") as f:
+                            f.write(f"{file.stem_uid}\n")
+                        continue
                 else:
                     print("Warning: couldn't check unicity from is_primary_data column")
+                if backed.obs.assay_ontology_term_id[0] in self.assays_to_drop:
+                    print(f"{file.key} is in the assay drop list.. dropping")
+                    continue
                 if backed.shape[1] < 1000:
                     print(
                         f"{file.key} only contains less than 1000 genes and is likely not scRNAseq... dropping"
@@ -556,37 +557,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 +622,7 @@ class LaminPreprocessor(Preprocessor):
                         continue
                     else:
                         raise e
-
+                gc.collect()
                 # issues with KLlggfw6I6lvmbqiZm46
             if self.keep_files:
                 # Reconstruct collection using keys
@@ -716,7 +732,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 +741,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 +750,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 +760,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 +789,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 +808,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 +881,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 +899,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 +915,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(