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

scdataloader/__init__.py

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

scdataloader/__main__.py

@@ -0,0 +1,188 @@
+import argparse
+from scdataloader.preprocess import LaminPreprocessor
+import lamindb as ln
+from typing import Optional, Union
+
+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, required=True, 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(
+        "--keep_norm_layer",
+        type=bool,
+        default=False,
+        help="Determines whether to keep the normalization layer."
+    )
+    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=200,
+        help="Specifies the minimum non-zero genes."
+    )
+    parser.add_argument(
+        "--maxdropamount",
+        type=int,
+        default=2,
+        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."
+    )
+    args = parser.parse_args()
+
+    # Load the collection
+    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()
+
+    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,
+        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,
+    )
+
+    # Preprocess the dataset
+    preprocessed_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,263 @@
+import numpy as np
+from .utils import load_genes
+from torch import Tensor
+
+# class SimpleCollator:
+
+
+class Collator:
+    def __init__(
+        self,
+        organisms: list,
+        org_to_id: dict = None,
+        valid_genes: list = [],
+        max_len=2000,
+        n_bins=0,
+        add_zero_genes=200,
+        logp1=False,
+        norm_to=None,
+        how="all",
+        tp_name=None,
+        organism_name="organism_ontology_term_id",
+        class_names=[],
+    ):
+        """
+        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.
+
+        Args:
+            organisms (list): List of organisms to be considered for gene expression data.
+            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.
+            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
+        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)
+        )
+        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)
+        for organism in set(self.genedf.organism):
+            ogenedf = self.genedf[self.genedf.organism == organism]
+            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]}
+            )
+            if len(valid_genes) > 0:
+                self.accepted_genes.update({org: ogenedf.index.isin(valid_genes)})
+
+    def __call__(self, batch):
+        """
+        __call__ is a special method in Python that is called when an instance of the class is called.
+
+        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 = []
+        for elem in batch:
+            organism_id = elem[self.organism_name]
+            if organism_id not in self.organism_ids:
+                continue
+            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]
+            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)
+                ]
+            elif self.how == "all":
+                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":
+                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])
+
+            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)
+
+        # 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
+        return {
+            "x": Tensor(expr),
+            "genes": Tensor(gene_locs).int(),
+            "class": Tensor(other_classes).int(),
+            "tp": Tensor(tp),
+            "depth": Tensor(total_count),
+        }
+
+
+class AnnDataCollator(Collator):
+    def __init__(self, *args, **kwargs):
+        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 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):
+        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):
+    def __call__(self, batch):
+        super().__call__(batch)
+        # binning
+
+        # tokenize the empty locations
+
+
+class scPRINTCollator(Collator):
+    def __call__(self, batch):
+        super().__call__(batch)