scdataloader 0.0.3__py3-none-any.whl → 1.0.1__py3-none-any.whl

scdataloader/preprocess.py

@@ -9,6 +9,7 @@ import scanpy as sc
 from anndata import AnnData
 from django.db import IntegrityError
 from scipy.sparse import csr_matrix
+import os
 
 from scdataloader import utils as data_utils
 
@@ -18,6 +19,8 @@ FULL_LENGTH_ASSAYS = [
     "EFO:0008931",
 ]
 
+MAXFILESIZE = 10_000_000_000
+
 
 class Preprocessor:
     """
@@ -30,23 +33,27 @@ class Preprocessor:
         filter_gene_by_counts: Union[int, bool] = False,
         filter_cell_by_counts: Union[int, bool] = False,
         normalize_sum: float = 1e4,
-        keep_norm_layer: bool = False,
         subset_hvg: int = 0,
+        use_layer: Optional[str] = None,
+        is_symbol: bool = False,
         hvg_flavor: str = "seurat_v3",
         binning: Optional[int] = None,
         result_binned_key: str = "X_binned",
         length_normalize: bool = False,
-        force_preprocess=False,
-        min_dataset_size=100,
-        min_valid_genes_id=10_000,
-        min_nnz_genes=200,
-        maxdropamount=2,
-        madoutlier=5,
-        pct_mt_outlier=8,
-        batch_key=None,
-        skip_validate=False,
+        force_preprocess: bool = False,
+        min_dataset_size: int = 100,
+        min_valid_genes_id: int = 10_000,
+        min_nnz_genes: int = 200,
+        maxdropamount: int = 50,
+        madoutlier: int = 5,
+        pct_mt_outlier: int = 8,
+        batch_key: Optional[str] = None,
+        skip_validate: bool = False,
         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"],
+        use_raw: bool = True,
     ) -> None:
         """
         Initializes the preprocessor and configures the workflow steps.
@@ -67,14 +74,34 @@ class Preprocessor:
             binning (int, optional): Determines whether to bin the data into discrete values of number of bins provided.
             result_binned_key (str, optional): Specifies the key of :class:`~anndata.AnnData` to store the binned data.
                 Defaults to "X_binned".
+            length_normalize (bool, optional): Determines whether to length normalize the data.
+                Defaults to False.
+            force_preprocess (bool, optional): Determines whether to bypass the check of raw counts.
+                Defaults to False.
+            min_dataset_size (int, optional): The minimum size required for a dataset to be kept.
+                Defaults to 100.
+            min_valid_genes_id (int, optional): The minimum number of valid genes to keep a dataset.
+                Defaults to 10_000.
+            min_nnz_genes (int, optional): The minimum number of non-zero genes to keep a cell.
+                Defaults to 200.
+            maxdropamount (int, optional): The maximum amount of dropped cells per dataset. (2 for 50% drop, 3 for 33% drop, etc.)
+                Defaults to 2.
+            madoutlier (int, optional): The maximum absolute deviation of the outlier samples.
+                Defaults to 5.
+            pct_mt_outlier (int, optional): The maximum percentage of mitochondrial genes outlier.
+                Defaults to 8.
+            batch_key (str, optional): The key of :class:`~anndata.AnnData.obs` to use for batch information.
+                This arg is used in the highly variable gene selection step.
+            skip_validate (bool, optional): Determines whether to skip the validation step.
+                Defaults to False.
         """
         self.filter_gene_by_counts = filter_gene_by_counts
         self.filter_cell_by_counts = filter_cell_by_counts
         self.normalize_sum = normalize_sum
-        self.keep_norm_layer = keep_norm_layer
         self.subset_hvg = subset_hvg
         self.hvg_flavor = hvg_flavor
         self.binning = binning
+        self.organisms = organisms
         self.result_binned_key = result_binned_key
         self.additional_preprocess = additional_preprocess
         self.additional_postprocess = additional_postprocess
@@ -88,45 +115,71 @@ class Preprocessor:
         self.batch_key = batch_key
         self.length_normalize = length_normalize
         self.skip_validate = skip_validate
+        self.use_layer = use_layer
+        self.is_symbol = is_symbol
+        self.do_postp = do_postp
+        self.use_raw = use_raw
 
     def __call__(self, adata) -> AnnData:
+        if adata[0].obs.organism_ontology_term_id.iloc[0] not in self.organisms:
+            raise ValueError(
+                "we cannot work with this organism",
+                adata[0].obs.organism_ontology_term_id.iloc[0],
+            )
         if self.additional_preprocess is not None:
             adata = self.additional_preprocess(adata)
-        if adata.raw is not None:
+        if adata.raw is not None and self.use_raw:
             adata.X = adata.raw.X
             del adata.raw
+        if self.use_layer is not None:
+            adata.X = adata.layers[self.use_layer]
         if adata.layers is not None:
+            if "counts" in adata.layers.keys():
+                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())
             del adata.layers
         if len(adata.varm.keys()) > 0:
             del adata.varm
-        if len(adata.obsm.keys()) > 0:
+        if len(adata.obsm.keys()) > 0 and self.do_postp:
             del adata.obsm
-        if len(adata.obsp.keys()) > 0:
+        if len(adata.obsp.keys()) > 0 and self.do_postp:
             del adata.obsp
         if len(adata.uns.keys()) > 0:
             del adata.uns
         if len(adata.varp.keys()) > 0:
             del adata.varp
         # check that it is a count
-        if (
-            np.abs(adata.X.astype(int) - adata.X).sum() and not self.force_preprocess
-        ):  # check if likely raw data
-            raise ValueError(
-                "Data is not raw counts, please check layers, find raw data, or bypass with force_preprocess"
-            )
+        print("checking raw counts")
+        if np.abs(
+            adata[:50_000].X.astype(int) - adata[:50_000].X
+        ).sum():  # check if likely raw data
+            if not self.force_preprocess:
+                raise ValueError(
+                    "Data is not raw counts, please check layers, find raw data, or bypass with force_preprocess"
+                )
+            else:
+                print(
+                    "Data is not raw counts, please check layers, find raw data, or bypass with force_preprocess"
+                )
             # please check layers
             # if not available count drop
+        prevsize = adata.shape[0]
+        # dropping non primary
+        if "is_primary_data" in adata.obs.columns:
+            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")
+        print(
+            "removed {} non primary cells, {} renamining".format(
+                prevsize - adata.shape[0], adata.shape[0]
+            )
+        )
         # # cleanup and dropping low expressed genes and unexpressed cells
         prevsize = adata.shape[0]
         adata.obs["nnz"] = np.array(np.sum(adata.X != 0, axis=1).flatten())[0]
-        adata = adata[
-            (adata.obs["nnz"] > self.min_nnz_genes)
-            # or if slide-seq
-            | (
-                (adata.obs.assay_ontology_term_id == "EFO:0030062")
-                & (adata.obs["nnz"] > (self.min_nnz_genes / 3))
-            )
-        ]
+        adata = adata[(adata.obs["nnz"] > self.min_nnz_genes)]
         if self.filter_gene_by_counts:
             sc.pp.filter_genes(adata, min_counts=self.filter_gene_by_counts)
         if self.filter_cell_by_counts:
@@ -145,12 +198,29 @@ class Preprocessor:
                 "Dataset dropped due to low expressed genes and unexpressed cells: current size: "
                 + str(adata.shape[0])
             )
-        # dropping non primary
-        adata = adata[adata.obs.is_primary_data]
-        if adata.shape[0] < self.min_dataset_size:
-            raise ValueError(
-                "Dataset dropped because contains too many secondary cells"
+        print(
+            "filtered out {} cells, {} renamining".format(
+                prevsize - adata.shape[0], adata.shape[0]
+            )
+        )
+
+        if self.is_symbol or not adata.var.index.str.contains("ENSG").any():
+            if not adata.var.index.str.contains("ENSG").any():
+                print("No ENSG genes found, assuming gene symbols...")
+            genesdf["ensembl_gene_id"] = genesdf.index
+            var = (
+                adata.var.merge(
+                    genesdf.drop_duplicates("symbol").set_index("symbol", drop=False),
+                    left_index=True,
+                    right_index=True,
+                    how="inner",
+                )
+                .sort_values(by="ensembl_gene_id")
+                .set_index("ensembl_gene_id")
             )
+            adata = adata[:, var["symbol"]]
+            adata.var = var
+            genesdf = genesdf.set_index("ensembl_gene_id")
 
         intersect_genes = set(adata.var.index).intersection(set(genesdf.index))
         print(f"Removed {len(adata.var.index) - len(intersect_genes)} genes.")
@@ -169,36 +239,39 @@ class Preprocessor:
         # do a validation function
         adata.uns["unseen_genes"] = list(unseen)
         if not self.skip_validate:
+            print("validating")
             data_utils.validate(adata, organism=adata.obs.organism_ontology_term_id[0])
-        # length normalization
-        if (
-            adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS).any()
-            and self.length_normalize
-        ):
-            subadata = data_utils.length_normalize(
-                adata[adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS)],
-            )
+            # length normalization
+            if (
+                adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS).any()
+                and self.length_normalize
+            ):
+                print("doing length norm")
+                subadata = data_utils.length_normalize(
+                    adata[adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS)],
+                )
 
-            adata = ad.concat(
-                [
-                    adata[
-                        ~adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS)
+                adata = ad.concat(
+                    [
+                        adata[
+                            ~adata.obs["assay_ontology_term_id"].isin(
+                                FULL_LENGTH_ASSAYS
+                            )
+                        ],
+                        subadata,
                     ],
-                    subadata,
-                ],
-                axis=0,
-                join="outer",
-                merge="only",
-            )
-        # step 3: normalize total
-        adata.layers["clean"] = sc.pp.log1p(
-            sc.pp.normalize_total(adata, target_sum=self.normalize_sum, inplace=False)[
-                "X"
-            ]
-        )
+                    axis=0,
+                    join="outer",
+                    merge="only",
+                )
 
         # QC
+
         adata.var[genesdf.columns] = genesdf.loc[adata.var.index]
+        for name in ["stable_id", "created_at", "updated_at"]:
+            if name in adata.var.columns:
+                adata.var = adata.var.drop(columns=name)
+        print("startin QC")
         sc.pp.calculate_qc_metrics(
             adata, qc_vars=["mt", "ribo", "hb"], inplace=True, percent_top=[20]
         )
@@ -224,31 +297,38 @@ class Preprocessor:
         #    raise Exception("More than 50% of the dataset has been dropped due to outliers.")
         # adata = adata[(~adata.obs.outlier) & (~adata.obs.mt_outlier)].copy()
         # remaining
-        # step 5: subset hvg
-        if self.subset_hvg:
-            sc.pp.highly_variable_genes(
-                adata,
-                layer="clean",
-                n_top_genes=self.subset_hvg,
-                batch_key=self.batch_key,
-                flavor=self.hvg_flavor,
-                subset=False,
-            )
+
         # based on the topometry paper https://www.biorxiv.org/content/10.1101/2022.03.14.484134v2
         # https://rapids-singlecell.readthedocs.io/en/latest/api/generated/rapids_singlecell.pp.pca.html#rapids_singlecell.pp.pca
-
-        adata.obsm["clean_pca"] = sc.pp.pca(
-            adata.layers["clean"],
-            n_comps=300 if adata.shape[0] > 300 else adata.shape[0] - 2,
-        )
-        sc.pp.neighbors(adata, use_rep="clean_pca")
-        sc.tl.leiden(adata, key_added="leiden_3", resolution=3.0)
-        sc.tl.leiden(adata, key_added="leiden_2", resolution=2.0)
-        sc.tl.leiden(adata, key_added="leiden_1", resolution=1.0)
-        sc.tl.umap(adata)
-        # additional
-        if self.additional_postprocess is not None:
-            adata = self.additional_postprocess(adata)
+        if self.do_postp:
+            print("normalize")
+            adata.layers["clean"] = sc.pp.log1p(
+                sc.pp.normalize_total(
+                    adata, target_sum=self.normalize_sum, inplace=False
+                )["X"]
+            )
+            # step 5: subset hvg
+            if self.subset_hvg:
+                sc.pp.highly_variable_genes(
+                    adata,
+                    layer="clean",
+                    n_top_genes=self.subset_hvg,
+                    batch_key=self.batch_key,
+                    flavor=self.hvg_flavor,
+                    subset=False,
+                )
+            adata.obsm["clean_pca"] = sc.pp.pca(
+                adata.layers["clean"],
+                n_comps=300 if adata.shape[0] > 300 else adata.shape[0] - 2,
+            )
+            sc.pp.neighbors(adata, use_rep="clean_pca")
+            sc.tl.leiden(adata, key_added="leiden_3", resolution=3.0)
+            sc.tl.leiden(adata, key_added="leiden_2", resolution=2.0)
+            sc.tl.leiden(adata, key_added="leiden_1", resolution=1.0)
+            sc.tl.umap(adata)
+            # additional
+            if self.additional_postprocess is not None:
+                adata = self.additional_postprocess(adata)
         adata = adata[:, adata.var.sort_index().index]
         # create random ids for all cells
         adata.obs.index = [str(uuid4()) for _ in range(adata.shape[0])]
@@ -296,6 +376,7 @@ class Preprocessor:
                 bin_edges.append(np.concatenate([[0], bins]))
             adata.layers[self.result_binned_key] = np.stack(binned_rows)
             adata.obsm["bin_edges"] = np.stack(bin_edges)
+        print("done")
         return adata
 
 
@@ -306,12 +387,14 @@ class LaminPreprocessor(Preprocessor):
         erase_prev_dataset: bool = False,
         cache: bool = True,
         stream: bool = False,
+        keep_files: bool = True,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.erase_prev_dataset = erase_prev_dataset
         self.cache = cache
         self.stream = stream
+        self.keep_files = keep_files
 
     def __call__(
         self,
@@ -319,7 +402,7 @@ class LaminPreprocessor(Preprocessor):
         name="preprocessed dataset",
         description="preprocessed dataset using scprint",
         start_at=0,
-        version="2",
+        version=2,
     ):
         """
         format controls the different input value wrapping, including categorical
@@ -334,49 +417,97 @@ class LaminPreprocessor(Preprocessor):
         all_ready_processed_keys = set()
         if self.cache:
             for i in ln.Artifact.filter(description=description):
-                all_ready_processed_keys.add(i.initial_version.key)
+                all_ready_processed_keys.add(i.stem_uid)
         if isinstance(data, AnnData):
-            return self.preprocess(data)
+            return super().__call__(data)
         elif isinstance(data, ln.Collection):
             for i, file in enumerate(data.artifacts.all()[start_at:]):
                 # use the counts matrix
                 print(i)
-                if file.key in all_ready_processed_keys:
-                    print(f"{file.key} is already processed")
+                if file.stem_uid in all_ready_processed_keys:
+                    print(f"{file.stem_uid} is already processed... not preprocessing")
                     continue
                 print(file)
-                if file.backed().obs.is_primary_data.sum() == 0:
-                    print(f"{file.key} only contains non primary cells")
+                backed = file.backed()
+                if backed.obs.is_primary_data.sum() == 0:
+                    print(f"{file.key} only contains non primary cells.. dropping")
+                    continue
+                if backed.shape[1] < 1000:
+                    print(
+                        f"{file.key} only contains less than 1000 genes and is likely not scRNAseq... dropping"
+                    )
                     continue
-                adata = file.load(stream=self.stream)
+                if file.size <= MAXFILESIZE:
+                    adata = file.load(stream=self.stream)
+                    print(adata)
+                else:
+                    badata = backed
+                    print(badata)
 
-                print(adata)
                 try:
-                    adata = super().__call__(adata)
+                    if file.size > MAXFILESIZE:
+                        print(
+                            f"dividing the dataset as it is too large: {file.size//1_000_000_000}Gb"
+                        )
+                        num_blocks = int(np.ceil(file.size / (MAXFILESIZE / 2)))
+                        block_size = int(
+                            (np.ceil(badata.shape[0] / 30_000) * 30_000) // num_blocks
+                        )
+                        print("num blocks ", num_blocks)
+                        for i in range(num_blocks):
+                            start_index = i * block_size
+                            end_index = min((i + 1) * block_size, badata.shape[0])
+                            block = badata[start_index:end_index].to_memory()
+                            print(block)
+                            block = super().__call__(block)
+                            myfile = ln.Artifact(
+                                block,
+                                is_new_version_of=file,
+                                description=description,
+                                version=str(version) + "_s" + str(i),
+                            )
+                            myfile.save()
+                            if self.keep_files:
+                                files.append(myfile)
+                            else:
+                                del myfile
+                                del block
+
+                    else:
+                        adata = super().__call__(adata)
+                        myfile = ln.Artifact(
+                            adata,
+                            is_new_version_of=file,
+                            description=description,
+                            version=str(version),
+                        )
+                        myfile.save()
+                        if self.keep_files:
+                            files.append(myfile)
+                        else:
+                            del myfile
+                            del adata
 
                 except ValueError as v:
-                    if v.args[0].startswith(
-                        "Dataset dropped because contains too many secondary"
-                    ):
+                    if v.args[0].startswith("we cannot work with this organism"):
                         print(v)
                         continue
                     else:
                         raise v
-                for name in ["stable_id", "created_at", "updated_at"]:
-                    if name in adata.var.columns:
-                        adata.var = adata.var.drop(columns=name)
-                myfile = ln.Artifact(
-                    adata,
-                    is_new_version_of=file,
-                    description=description,
-                    version=version,
-                )
+                except Exception as e:
+                    if e.args[0].startswith("Dataset dropped due to"):
+                        print(e)
+                        continue
+                    else:
+                        raise e
+
                 # issues with KLlggfw6I6lvmbqiZm46
-                myfile.save()
-                files.append(myfile)
-            dataset = ln.Collection(files, name=name, description=description)
-            dataset.save()
-            return dataset
+            if self.keep_files:
+                dataset = ln.Collection(files, name=name, description=description)
+                dataset.save()
+                return dataset
+            else:
+                return
         else:
             raise ValueError("Please provide either anndata or ln.Collection")
 
@@ -498,7 +629,7 @@ def additional_preprocess(adata):
         ].astype(str)
         adata.obs.loc[loc, "tissue_ontology_term_id"] = adata.obs.loc[
             loc, "tissue_ontology_term_id"
-        ].str.replace(r" \(cell culture\)", "")
+        ].str.replace(" (cell culture)", "")
 
     loc = adata.obs["tissue_ontology_term_id"].str.contains("(organoid)", regex=False)
     if loc.sum() > 0:
@@ -508,7 +639,7 @@ def additional_preprocess(adata):
         ].astype(str)
         adata.obs.loc[loc, "tissue_ontology_term_id"] = adata.obs.loc[
             loc, "tissue_ontology_term_id"
-        ].str.replace(r" \(organoid\)", "")
+        ].str.replace(" (organoid)", "")
 
     loc = adata.obs["tissue_ontology_term_id"].str.contains("CL:", regex=False)
     if loc.sum() > 0: