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

scdataloader/VERSION

@@ -1 +1 @@
-
+0.7.0

scdataloader/__init__.py

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

scdataloader/__main__.py

@@ -0,0 +1,209 @@
+import argparse
+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."
+    )
+    parser.add_argument(
+        "--name", type=str, required=True, help="Name of the input dataset"
+    )
+    parser.add_argument(
+        "--new_name",
+        type=str,
+        default="preprocessed dataset",
+        help="Name of the preprocessed dataset.",
+    )
+    parser.add_argument(
+        "--description",
+        type=str,
+        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.",
+    )
+    parser.add_argument(
+        "--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."
+    )
+    parser.add_argument(
+        "--filter_gene_by_counts",
+        type=Union[int, bool],
+        default=False,
+        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.",
+    )
+    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(
+        "--subset_hvg",
+        type=int,
+        default=0,
+        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.",
+    )
+    parser.add_argument(
+        "--binning",
+        type=Optional[int],
+        default=None,
+        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.",
+    )
+    parser.add_argument(
+        "--length_normalize",
+        type=bool,
+        default=False,
+        help="Determines whether to normalize the length.",
+    )
+    parser.add_argument(
+        "--force_preprocess",
+        type=bool,
+        default=False,
+        help="Determines whether to force preprocessing.",
+    )
+    parser.add_argument(
+        "--min_dataset_size",
+        type=int,
+        default=100,
+        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.",
+    )
+    parser.add_argument(
+        "--min_nnz_genes",
+        type=int,
+        default=400,
+        help="Specifies the minimum non-zero genes.",
+    )
+    parser.add_argument(
+        "--maxdropamount",
+        type=int,
+        default=50,
+        help="Specifies the maximum drop amount.",
+    )
+    parser.add_argument(
+        "--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.",
+    )
+    parser.add_argument(
+        "--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.",
+    )
+    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()
+        )
+    else:
+        collection = ln.Collection.filter(name=args.name, version=args.version).first()
+
+    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,
+        subset_hvg=args.subset_hvg,
+        hvg_flavor=args.hvg_flavor,
+        binning=args.binning,
+        result_binned_key=args.result_binned_key,
+        length_normalize=args.length_normalize,
+        force_preprocess=args.force_preprocess,
+        min_dataset_size=args.min_dataset_size,
+        min_valid_genes_id=args.min_valid_genes_id,
+        min_nnz_genes=args.min_nnz_genes,
+        maxdropamount=args.maxdropamount,
+        madoutlier=args.madoutlier,
+        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
+    preprocessor(
+        collection,
+        name=args.new_name,
+        description=args.description,
+        start_at=args.start_at,
+        version=args.new_version,
+    )
+
+    print(
+        f"Preprocessed dataset saved with version {args.version} and name {args.new_name}."
+    )
+
+
+if __name__ == "__main__":
+    main()

scdataloader/collator.py

@@ -0,0 +1,307 @@
+import numpy as np
+from .utils import load_genes
+from torch import Tensor, long
+
+# class SimpleCollator:
+
+
+class Collator:
+    def __init__(
+        self,
+        organisms: list,
+        how="all",
+        org_to_id: dict = None,
+        valid_genes: list = [],
+        max_len=2000,
+        add_zero_genes=0,
+        logp1=False,
+        norm_to=None,
+        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
+        organization and preparation of gene expression data from different organisms,
+        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.
+            valid_genes (list, optional): List of genes from the datasets, to be considered. Defaults to [].
+                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.
+        """
+        self.organisms = organisms
+        self.max_len = max_len
+        self.n_bins = n_bins
+        self.add_zero_genes = add_zero_genes
+        self.logp1 = logp1
+        self.norm_to = norm_to
+        self.org_to_id = org_to_id
+        self.how = how
+        self.organism_ids = (
+            set([org_to_id[k] for k in organisms])
+            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
+
+        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(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__ applies the collator to a minibatch of data
+
+        Args:
+            batch (list[dict[str: array]]): List of dicts of arrays containing gene expression data.
+                the first list is for the different samples, the second list is for the different elements with
+                elem["x"]: gene expression
+                elem["organism_name"]: organism ontology term id
+                elem["tp_name"]: heat diff
+                elem["class_names.."]: other classes
+
+        Returns:
+            list[Tensor]: List of tensors containing the collated data.
+        """
+        # do count selection
+        # get the unseen info and don't add any unseen
+        # get the I most expressed genes, add randomly some unexpressed genes that are not unseen
+        exprs = []
+        total_count = []
+        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":
+                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 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 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) or (self.max_len > len(nnz_loc))
+            ) and self.how not in ["all", "some"]:
+                zero_loc = np.where(expr == 0)[0]
+                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)
+            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])
+            else:
+                tp.append(0)
+
+            other_classes.append([elem[i] for i in self.class_names])
+
+        expr = np.array(exprs)
+        tp = np.array(tp)
+        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:
+            expr = (expr * self.norm_to) / total_count[:, None]
+        if self.logp1:
+            expr = np.log2(1 + expr)
+
+        # do binning of counts
+        if self.n_bins:
+            pass
+
+        # find the associated gene ids (given the species)
+
+        # get the NN cells
+
+        # do encoding / selection a la scGPT
+
+        # do encoding of graph location
+        # encode all the edges in some sparse way
+        # normalizing total counts between 0,1
+        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):
+        exprs = []
+        total_count = []
+        other_classes = []
+        gene_locs = []
+        tp = []
+        for elem in batch:
+            organism_id = elem.obs[self.organism_name]
+            if organism_id.item() not in self.organism_ids:
+                print(organism_id)
+            expr = np.array(elem.X[0])
+
+            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]
+            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)
+                ]
+            else:
+                raise ValueError("how must be either most expr or random expr")
+            if self.add_zero_genes > 0:
+                zero_loc = np.where(expr == 0)[0]
+                zero_loc = [
+                    np.random.choice(len(zero_loc), self.add_zero_genes, replace=False)
+                ]
+                loc = np.concatenate((loc, zero_loc), axis=None)
+            exprs.append(expr[loc])
+            gene_locs.append(loc + self.start_idx[organism_id.item()])
+
+            if self.tp_name is not None:
+                tp.append(elem.obs[self.tp_name])
+            else:
+                tp.append(0)
+
+            other_classes.append([elem.obs[i].values[0] for i in self.class_names])
+
+        expr = np.array(exprs)
+        tp = np.array(tp)
+        gene_locs = np.array(gene_locs)
+        total_count = np.array(total_count)
+        other_classes = np.array(other_classes)
+        return {
+            "x": Tensor(expr),
+            "genes": Tensor(gene_locs).int(),
+            "depth": Tensor(total_count),
+            "class": Tensor(other_classes),
+        }
+
+
+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
+
+    def __call__(self, batch):
+        super().__call__(batch)
+        # normlization per gene
+
+        # tokenize the empty locations
+
+
+class scGPTCollator(Collator):
+    """
+    scGPTCollator to finish
+    """
+
+    def __call__(self, batch):
+        super().__call__(batch)
+        # binning
+
+        # tokenize the empty locations
+
+
+class scPRINTCollator(Collator):
+    def __call__(self, batch):
+        super().__call__(batch)

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": "",
+}