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

scdataloader/VERSION

@@ -1 +1 @@
-
+0.7.0

scdataloader/__init__.py

@@ -1,4 +1,4 @@
 from .data import Dataset
-from .dataloader import DataModule
+from .datamodule import DataModule
 from .preprocess import Preprocessor
 from .collator import *

scdataloader/__main__.py

@@ -1,8 +1,14 @@
 import argparse
-from scdataloader.preprocess import LaminPreprocessor
+from scdataloader.preprocess import (
+    LaminPreprocessor,
+    additional_preprocess,
+    additional_postprocess,
+)
 import lamindb as ln
 from typing import Optional, Union
 
+
+# scdataloader --instance="laminlabs/cellxgene" --name="cellxgene-census" --version="2023-12-15" --description="preprocessed for scprint" --new_name="scprint main" --start_at=39
 def main():
     parser = argparse.ArgumentParser(
         description="Preprocess datasets in a given lamindb collection."
@@ -11,22 +17,31 @@ def main():
         "--name", type=str, required=True, help="Name of the input dataset"
     )
     parser.add_argument(
-        "--new_name", type=str, required=True, help="Name of the preprocessed dataset."
+        "--new_name",
+        type=str,
+        default="preprocessed dataset",
+        help="Name of the preprocessed dataset.",
     )
     parser.add_argument(
         "--description",
         type=str,
-        default="preprocessed by scDataLoader"
+        default="preprocessed by scDataLoader",
         help="Description of the preprocessed dataset.",
     )
     parser.add_argument(
         "--start_at", type=int, default=0, help="Position to start preprocessing at."
     )
     parser.add_argument(
-        "--new_version", type=str, default="2", help="Version of the output dataset and files."
+        "--new_version",
+        type=str,
+        default="2",
+        help="Version of the output dataset and files.",
     )
     parser.add_argument(
-        "--instance", type=str, default=None, help="Instance storing the input dataset, if not local"
+        "--instance",
+        type=str,
+        default=None,
+        help="Instance storing the input dataset, if not local",
     )
     parser.add_argument(
         "--version", type=str, default=None, help="Version of the input dataset."
@@ -35,125 +50,127 @@ def main():
         "--filter_gene_by_counts",
         type=Union[int, bool],
         default=False,
-        help="Determines whether to filter genes by counts."
+        help="Determines whether to filter genes by counts.",
     )
     parser.add_argument(
         "--filter_cell_by_counts",
         type=Union[int, bool],
         default=False,
-        help="Determines whether to filter cells by counts."
+        help="Determines whether to filter cells by counts.",
     )
     parser.add_argument(
         "--normalize_sum",
         type=float,
         default=1e4,
-        help="Determines whether to normalize the total counts of each cell to a specific value."
-    )
-    parser.add_argument(
-        "--keep_norm_layer",
-        type=bool,
-        default=False,
-        help="Determines whether to keep the normalization layer."
+        help="Determines whether to normalize the total counts of each cell to a specific value.",
     )
     parser.add_argument(
         "--subset_hvg",
         type=int,
         default=0,
-        help="Determines whether to subset highly variable genes."
+        help="Determines whether to subset highly variable genes.",
     )
     parser.add_argument(
         "--hvg_flavor",
         type=str,
         default="seurat_v3",
-        help="Specifies the flavor of highly variable genes selection."
+        help="Specifies the flavor of highly variable genes selection.",
     )
     parser.add_argument(
         "--binning",
         type=Optional[int],
         default=None,
-        help="Determines whether to bin the data into discrete values of number of bins provided."
+        help="Determines whether to bin the data into discrete values of number of bins provided.",
     )
     parser.add_argument(
         "--result_binned_key",
         type=str,
         default="X_binned",
-        help="Specifies the key of AnnData to store the binned data."
+        help="Specifies the key of AnnData to store the binned data.",
     )
     parser.add_argument(
         "--length_normalize",
         type=bool,
         default=False,
-        help="Determines whether to normalize the length."
+        help="Determines whether to normalize the length.",
     )
     parser.add_argument(
         "--force_preprocess",
         type=bool,
         default=False,
-        help="Determines whether to force preprocessing."
+        help="Determines whether to force preprocessing.",
     )
     parser.add_argument(
         "--min_dataset_size",
         type=int,
         default=100,
-        help="Specifies the minimum dataset size."
+        help="Specifies the minimum dataset size.",
     )
     parser.add_argument(
         "--min_valid_genes_id",
         type=int,
         default=10_000,
-        help="Specifies the minimum valid genes id."
+        help="Specifies the minimum valid genes id.",
     )
     parser.add_argument(
         "--min_nnz_genes",
         type=int,
-        default=200,
-        help="Specifies the minimum non-zero genes."
+        default=400,
+        help="Specifies the minimum non-zero genes.",
     )
     parser.add_argument(
         "--maxdropamount",
         type=int,
-        default=2,
-        help="Specifies the maximum drop amount."
+        default=50,
+        help="Specifies the maximum drop amount.",
     )
     parser.add_argument(
-        "--madoutlier",
-        type=int,
-        default=5,
-        help="Specifies the MAD outlier."
+        "--madoutlier", type=int, default=5, help="Specifies the MAD outlier."
     )
     parser.add_argument(
         "--pct_mt_outlier",
         type=int,
         default=8,
-        help="Specifies the percentage of MT outlier."
+        help="Specifies the percentage of MT outlier.",
     )
     parser.add_argument(
-        "--batch_key",
-        type=Optional[str],
-        default=None,
-        help="Specifies the batch key."
+        "--batch_key", type=Optional[str], default=None, help="Specifies the batch key."
     )
     parser.add_argument(
         "--skip_validate",
         type=bool,
         default=False,
-        help="Determines whether to skip validation."
+        help="Determines whether to skip validation.",
+    )
+    parser.add_argument(
+        "--do_postp",
+        type=bool,
+        default=False,
+        help="Determines whether to do postprocessing.",
     )
     args = parser.parse_args()
 
     # Load the collection
+    # if not args.preprocess:
+    #    print("Only preprocess is available for now")
+    #    return
     if args.instance is not None:
-        collection = ln.Collection.using(instance=args.instance).filter(name=args.name, version=args.version).first()
-
-    collection = ln.Collection.filter(name=args.name, version=args.version).first()
+        collection = (
+            ln.Collection.using(instance=args.instance)
+            .filter(name=args.name, version=args.version)
+            .first()
+        )
+    else:
+        collection = ln.Collection.filter(name=args.name, version=args.version).first()
 
-    print("using the dataset ",collection, " of size ",len(collection.artifacts.all()))
+    print(
+        "using the dataset ", collection, " of size ", len(collection.artifacts.all())
+    )
     # Initialize the preprocessor
     preprocessor = LaminPreprocessor(
         filter_gene_by_counts=args.filter_gene_by_counts,
         filter_cell_by_counts=args.filter_cell_by_counts,
         normalize_sum=args.normalize_sum,
-        keep_norm_layer=args.keep_norm_layer,
         subset_hvg=args.subset_hvg,
         hvg_flavor=args.hvg_flavor,
         binning=args.binning,
@@ -168,10 +185,14 @@ def main():
         pct_mt_outlier=args.pct_mt_outlier,
         batch_key=args.batch_key,
         skip_validate=args.skip_validate,
+        do_postp=args.do_postp,
+        additional_preprocess=additional_preprocess,
+        additional_postprocess=additional_postprocess,
+        keep_files=False,
     )
 
     # Preprocess the dataset
-    preprocessed_dataset = preprocessor(
+    preprocessor(
         collection,
         name=args.new_name,
         description=args.description,

scdataloader/collator.py

@@ -1,6 +1,6 @@
 import numpy as np
 from .utils import load_genes
-from torch import Tensor
+from torch import Tensor, long
 
 # class SimpleCollator:
 
@@ -9,17 +9,18 @@ class Collator:
     def __init__(
         self,
         organisms: list,
+        how="all",
         org_to_id: dict = None,
         valid_genes: list = [],
         max_len=2000,
-        n_bins=0,
-        add_zero_genes=200,
+        add_zero_genes=0,
         logp1=False,
         norm_to=None,
-        how="all",
+        n_bins=0,
         tp_name=None,
         organism_name="organism_ontology_term_id",
         class_names=[],
+        genelist=[],
     ):
         """
         This class is responsible for collating data for the scPRINT model. It handles the
@@ -27,20 +28,29 @@ class Collator:
         allowing for various configurations such as maximum gene list length, normalization,
         and selection method for gene expression.
 
+        This Collator should work with scVI's dataloader as well!
+
         Args:
             organisms (list): List of organisms to be considered for gene expression data.
+                it will drop any other organism it sees (might lead to batches of different sizes!)
+            how (flag, optional): Method for selecting gene expression. Defaults to "most expr".
+                one of ["most expr", "random expr", "all", "some"]:
+                "most expr": selects the max_len most expressed genes,
+                if less genes are expressed, will sample random unexpressed genes,
+                "random expr": uses a random set of max_len expressed genes.
+                if less genes are expressed, will sample random unexpressed genes
+                "all": uses all genes
+                "some": uses only the genes provided through the genelist param
             org_to_id (dict): Dictionary mapping organisms to their respective IDs.
-            labels (list, optional): List of labels for the data. Defaults to [].
             valid_genes (list, optional): List of genes from the datasets, to be considered. Defaults to [].
-            max_len (int, optional): Maximum length of the gene list. Defaults to 2000.
-            n_bins (int, optional): Number of bins for binning the data. Defaults to 0.
-            add_zero_genes (int, optional): Number of zero genes to add. Defaults to 200.
+                it will drop any other genes from the input expression data (usefull when your model only works on some genes)
+            max_len (int, optional): Maximum number of genes to use (for random expr and most expr). Defaults to 2000.
+            n_bins (int, optional): Number of bins for binning the data. Defaults to 0. meaning, no binning of expression.
+            add_zero_genes (int, optional): Number of additional unexpressed genes to add to the input data. Defaults to 0.
             logp1 (bool, optional): If True, logp1 normalization is applied. Defaults to False.
             norm_to (str, optional): Normalization method to be applied. Defaults to None.
-            how (str, optional): Method for selecting gene expression. Defaults to "most expr".
         """
         self.organisms = organisms
-        self.valid_genes = valid_genes
         self.max_len = max_len
         self.n_bins = n_bins
         self.add_zero_genes = add_zero_genes
@@ -53,6 +63,8 @@ class Collator:
             if org_to_id is not None
             else set(organisms)
         )
+        if self.how == "some":
+            assert len(genelist) > 0, "if how is some, genelist must be provided"
         self.organism_name = organism_name
         self.tp_name = tp_name
         self.class_names = class_names
@@ -60,18 +72,21 @@ class Collator:
         self.start_idx = {}
         self.accepted_genes = {}
         self.genedf = load_genes(organisms)
+        self.to_subset = {}
         for organism in set(self.genedf.organism):
             ogenedf = self.genedf[self.genedf.organism == organism]
+            tot = self.genedf[self.genedf.index.isin(valid_genes)]
             org = org_to_id[organism] if org_to_id is not None else organism
-            self.start_idx.update(
-                {org: np.where(self.genedf.organism == organism)[0][0]}
-            )
+            self.start_idx.update({org: np.where(tot.organism == organism)[0][0]})
             if len(valid_genes) > 0:
                 self.accepted_genes.update({org: ogenedf.index.isin(valid_genes)})
+            if len(genelist) > 0:
+                df = ogenedf[ogenedf.index.isin(valid_genes)]
+                self.to_subset.update({org: df.index.isin(genelist)})
 
     def __call__(self, batch):
         """
-        __call__ is a special method in Python that is called when an instance of the class is called.
+        __call__ applies the collator to a minibatch of data
 
         Args:
             batch (list[dict[str: array]]): List of dicts of arrays containing gene expression data.
@@ -92,33 +107,62 @@ class Collator:
         other_classes = []
         gene_locs = []
         tp = []
+        dataset = []
+        nnz_loc = []
         for elem in batch:
             organism_id = elem[self.organism_name]
             if organism_id not in self.organism_ids:
                 continue
+            if "dataset" in elem:
+                dataset.append(elem["dataset"])
             expr = np.array(elem["x"])
             total_count.append(expr.sum())
             if len(self.accepted_genes) > 0:
                 expr = expr[self.accepted_genes[organism_id]]
             if self.how == "most expr":
-                loc = np.argsort(expr)[-(self.max_len) :][::-1]
+                nnz_loc = np.where(expr > 0)[0]
+                ma = self.max_len if self.max_len < len(nnz_loc) else len(nnz_loc)
+                loc = np.argsort(expr)[-(ma):][::-1]
+                # nnz_loc = [1] * 30_000
+                # loc = np.argsort(expr)[-(self.max_len) :][::-1]
             elif self.how == "random expr":
                 nnz_loc = np.where(expr > 0)[0]
                 loc = nnz_loc[
-                    np.random.choice(len(nnz_loc), self.max_len, replace=False)
+                    np.random.choice(
+                        len(nnz_loc),
+                        self.max_len if self.max_len < len(nnz_loc) else len(nnz_loc),
+                        replace=False,
+                        # p=(expr.max() + (expr[nnz_loc])*19) / expr.max(), # 20 at most times more likely to be selected
+                    )
                 ]
-            elif self.how == "all":
+            elif self.how in ["all", "some"]:
                 loc = np.arange(len(expr))
             else:
                 raise ValueError("how must be either most expr or random expr")
-            if self.add_zero_genes > 0 and self.how != "all":
+            if (
+                (self.add_zero_genes > 0) or (self.max_len > len(nnz_loc))
+            ) and self.how not in ["all", "some"]:
                 zero_loc = np.where(expr == 0)[0]
-                zero_loc = [
-                    np.random.choice(len(zero_loc), self.add_zero_genes, replace=False)
+                zero_loc = zero_loc[
+                    np.random.choice(
+                        len(zero_loc),
+                        self.add_zero_genes
+                        + (
+                            0
+                            if self.max_len < len(nnz_loc)
+                            else self.max_len - len(nnz_loc)
+                        ),
+                        replace=False,
+                    )
                 ]
                 loc = np.concatenate((loc, zero_loc), axis=None)
-            exprs.append(expr[loc])
-            gene_locs.append(loc + self.start_idx[organism_id])
+            expr = expr[loc]
+            loc = loc + self.start_idx[organism_id]
+            if self.how == "some":
+                expr = expr[self.to_subset[organism_id]]
+                loc = loc[self.to_subset[organism_id]]
+            exprs.append(expr)
+            gene_locs.append(loc)
 
             if self.tp_name is not None:
                 tp.append(elem[self.tp_name])
@@ -132,6 +176,7 @@ class Collator:
         gene_locs = np.array(gene_locs)
         total_count = np.array(total_count)
         other_classes = np.array(other_classes)
+        dataset = np.array(dataset)
 
         # normalize counts
         if self.norm_to is not None:
@@ -152,17 +197,26 @@ class Collator:
         # do encoding of graph location
         # encode all the edges in some sparse way
         # normalizing total counts between 0,1
-        return {
+        ret = {
             "x": Tensor(expr),
             "genes": Tensor(gene_locs).int(),
             "class": Tensor(other_classes).int(),
             "tp": Tensor(tp),
             "depth": Tensor(total_count),
         }
+        if len(dataset) > 0:
+            ret.update({"dataset": Tensor(dataset).to(long)})
+        return ret
 
 
 class AnnDataCollator(Collator):
     def __init__(self, *args, **kwargs):
+        """
+        AnnDataCollator Collator to use if working with AnnData's experimental dataloader (it is very slow!!!)
+
+        Args:
+            @see Collator
+        """
         super().__init__(*args, **kwargs)
 
     def __call__(self, batch):
@@ -218,28 +272,14 @@ class AnnDataCollator(Collator):
         }
 
 
-class SCVICollator(Collator):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def __call__(self, batch):
-        expr = batch["x"]
-        total_count = expr.sum(axis=1)
-        if self.how == "most expr":
-            loc = np.argsort(expr)[:, -(self.max_len) :][:, ::-1]
-        else:
-            raise ValueError("how must be either most expr or random expr")
-        if self.logp1:
-            expr = np.log2(1 + expr)
-        return {
-            "x": Tensor(expr[np.arange(expr.shape[0])[:, None], loc]),
-            "genes": Tensor(loc.copy()).int(),
-            "depth": Tensor(total_count),
-        }
-
-
 class GeneformerCollator(Collator):
     def __init__(self, *args, gene_norm_list: list, **kwargs):
+        """
+        GeneformerCollator to finish
+
+        Args:
+            gene_norm_list (list): the normalization of expression through all datasets, per gene.
+        """
         super().__init__(*args, **kwargs)
         self.gene_norm_list = gene_norm_list
 
@@ -251,6 +291,10 @@ class GeneformerCollator(Collator):
 
 
 class scGPTCollator(Collator):
+    """
+    scGPTCollator to finish
+    """
+
     def __call__(self, batch):
         super().__call__(batch)
         # binning

scdataloader/config.py

@@ -0,0 +1,106 @@
+LABELS_TOADD = {
+    "assay_ontology_term_id": {
+        "10x transcription profiling": "EFO:0030003",
+        "spatial transcriptomics": "EFO:0008994",
+        "10x 3' transcription profiling": "EFO:0030003",
+        "10x 5' transcription profiling": "EFO:0030004",
+    },
+    "disease_ontology_term_id": {
+        "metabolic disease": "MONDO:0005066",
+        "chronic kidney disease": "MONDO:0005300",
+        "chromosomal disorder": "MONDO:0019040",
+        "infectious disease": "MONDO:0005550",
+        "inflammatory disease": "MONDO:0021166",
+        # "immune system disease",
+        "disorder of development or morphogenesis": "MONDO:0021147",
+        "mitochondrial disease": "MONDO:0044970",
+        "psychiatric disorder": "MONDO:0002025",
+        "cancer or benign tumor": "MONDO:0002025",
+        "neoplasm": "MONDO:0005070",
+    },
+    "cell_type_ontology_term_id": {
+        "progenitor cell": "CL:0011026",
+        "hematopoietic cell": "CL:0000988",
+        "myoblast": "CL:0000056",
+        "myeloid cell": "CL:0000763",
+        "neuron": "CL:0000540",
+        "electrically active cell": "CL:0000211",
+        "epithelial cell": "CL:0000066",
+        "secretory cell": "CL:0000151",
+        "stem cell": "CL:0000034",
+        "non-terminally differentiated cell": "CL:0000055",
+        "supporting cell": "CL:0000630",
+    },
+}
+
+COARSE_TISSUE = {
+    "adipose tissue": "",
+    "bladder organ": "",
+    "blood": "",
+    "bone marrow": "",
+    "brain": "",
+    "breast": "",
+    "esophagus": "",
+    "eye": "",
+    "embryo": "",
+    "fallopian tube": "",
+    "gall bladder": "",
+    "heart": "",
+    "intestine": "",
+    "kidney": "",
+    "liver": "",
+    "lung": "",
+    "lymph node": "",
+    "musculature of body": "",
+    "nose": "",
+    "ovary": "",
+    "pancreas": "",
+    "placenta": "",
+    "skin of body": "",
+    "spinal cord": "",
+    "spleen": "",
+    "stomach": "",
+    "thymus": "",
+    "thyroid gland": "",
+    "tongue": "",
+    "uterus": "",
+}
+
+COARSE_ANCESTRY = {
+    "African": "",
+    "Chinese": "",
+    "East Asian": "",
+    "Eskimo": "",
+    "European": "",
+    "Greater Middle Eastern  (Middle Eastern, North African or Persian)": "",
+    "Hispanic or Latin American": "",
+    "Native American": "",
+    "Oceanian": "",
+    "South Asian": "",
+}
+
+COARSE_DEVELOPMENT_STAGE = {
+    "Embryonic human": "",
+    "Fetal": "",
+    "Immature": "",
+    "Mature": "",
+}
+
+COARSE_ASSAY = {
+    "10x 3'": "",
+    "10x 5'": "",
+    "10x multiome": "",
+    "CEL-seq2": "",
+    "Drop-seq": "",
+    "GEXSCOPE technology": "",
+    "inDrop": "",
+    "microwell-seq": "",
+    "sci-Plex": "",
+    "sci-RNA-seq": "",
+    "Seq-Well": "",
+    "Slide-seq": "",
+    "Smart-seq": "",
+    "SPLiT-seq": "",
+    "TruDrop": "",
+    "Visium Spatial Gene Expression": "",
+}