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

scdataloader/VERSION

@@ -1 +1 @@
-0.7.0
+1.0.5

scdataloader/__init__.py

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

scdataloader/__main__.py

@@ -10,6 +10,9 @@ 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():
+    """
+    main function to preprocess datasets in a given lamindb collection.
+    """
     parser = argparse.ArgumentParser(
         description="Preprocess datasets in a given lamindb collection."
     )

scdataloader/collator.py

@@ -1,26 +1,27 @@
 import numpy as np
-from .utils import load_genes
+from .utils import load_genes, downsample_profile
 from torch import Tensor, long
-
-# class SimpleCollator:
+from typing import Optional
 
 
 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=[],
+        organisms: list[str],
+        how: str = "all",
+        org_to_id: dict[str, int] = None,
+        valid_genes: list[str] = [],
+        max_len: int = 2000,
+        add_zero_genes: int = 0,
+        logp1: bool = False,
+        norm_to: Optional[float] = None,
+        n_bins: int = 0,
+        tp_name: Optional[str] = None,
+        organism_name: str = "organism_ontology_term_id",
+        class_names: list[str] = [],
+        genelist: list[str] = [],
+        downsample: Optional[float] = None,  # don't use it for training!
+        save_output: bool = False,
     ):
         """
         This class is responsible for collating data for the scPRINT model. It handles the
@@ -44,38 +45,57 @@ class Collator:
             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.
+            max_len (int, optional): Total 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.
+            norm_to (float, optional): Rescaling value of the normalization to be applied. Defaults to None.
+            organism_name (str, optional): Name of the organism ontology term id. Defaults to "organism_ontology_term_id".
+            tp_name (str, optional): Name of the heat diff. Defaults to None.
+            class_names (list, optional): List of other classes to be considered. Defaults to [].
+            genelist (list, optional): List of genes to be considered. Defaults to [].
+                If [] all genes will be considered
+            downsample (float, optional): Downsample the profile to a certain number of cells. Defaults to None.
+                This is usually done by the scPRINT model during training but this option allows you to do it directly from the collator
+            save_output (bool, optional): If True, saves the output to a file. Defaults to False.
+                This is mainly for debugging purposes
         """
         self.organisms = organisms
+        self.genedf = load_genes(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.save_output = save_output
         self.start_idx = {}
         self.accepted_genes = {}
-        self.genedf = load_genes(organisms)
+        self.downsample = downsample
         self.to_subset = {}
-        for organism in set(self.genedf.organism):
+        self._setup(org_to_id, valid_genes, genelist)
+
+    def _setup(self, org_to_id=None, valid_genes=[], genelist=[]):
+        self.org_to_id = org_to_id
+        self.to_subset = {}
+        self.accepted_genes = {}
+        self.start_idx = {}
+        self.organism_ids = (
+            set([org_to_id[k] for k in self.organisms])
+            if org_to_id is not None
+            else set(self.organisms)
+        )
+        for organism in self.organisms:
             ogenedf = self.genedf[self.genedf.organism == organism]
-            tot = self.genedf[self.genedf.index.isin(valid_genes)]
+            if len(valid_genes) > 0:
+                tot = self.genedf[self.genedf.index.isin(valid_genes)]
+            else:
+                tot = self.genedf
             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:
@@ -84,14 +104,14 @@ class Collator:
                 df = ogenedf[ogenedf.index.isin(valid_genes)]
                 self.to_subset.update({org: df.index.isin(genelist)})
 
-    def __call__(self, batch):
+    def __call__(self, batch) -> dict[str, Tensor]:
         """
         __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["X"]: gene expression
                 elem["organism_name"]: organism ontology term id
                 elem["tp_name"]: heat diff
                 elem["class_names.."]: other classes
@@ -113,9 +133,9 @@ class Collator:
             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"])
+            if "_storage_idx" in elem:
+                dataset.append(elem["_storage_idx"])
+            expr = np.array(elem["X"])
             total_count.append(expr.sum())
             if len(self.accepted_genes) > 0:
                 expr = expr[self.accepted_genes[organism_id]]
@@ -206,72 +226,17 @@ class Collator:
         }
         if len(dataset) > 0:
             ret.update({"dataset": Tensor(dataset).to(long)})
+        if self.downsample is not None:
+            ret["x"] = downsample_profile(ret["x"], self.downsample)
+        if self.save_output:
+            with open("collator_output.txt", "a") as f:
+                np.savetxt(f, ret["x"].numpy())
         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),
-        }
-
-
+#############
+#### WIP ####
+#############
 class GeneformerCollator(Collator):
     def __init__(self, *args, gene_norm_list: list, **kwargs):
         """

scdataloader/config.py

@@ -1,3 +1,9 @@
+"""
+Configuration file for scDataLoader
+
+Missing labels are added to the dataset to complete a better hierarchical tree
+"""
+
 LABELS_TOADD = {
     "assay_ontology_term_id": {
         "10x transcription profiling": "EFO:0030003",

scdataloader/data.py

@@ -1,14 +1,18 @@
 from dataclasses import dataclass, field
 
 import lamindb as ln
+
+# ln.connect("scprint")
+
 import bionty as bt
 import pandas as pd
 from torch.utils.data import Dataset as torchDataset
 from typing import Union, Optional, Literal
-from scdataloader import mapped
+from scdataloader.mapped import MappedCollection
 import warnings
 
 from anndata import AnnData
+from scipy.sparse import issparse
 
 from scdataloader.utils import get_ancestry_mapping, load_genes
 
@@ -58,30 +62,31 @@ class Dataset(torchDataset):
             "sex_ontology_term_id",
             #'dataset_id',
             #'cell_culture',
-            #"dpt_group",
-            #"heat_diff",
-            #"nnz",
+            # "dpt_group",
+            # "heat_diff",
+            # "nnz",
         ]
     )
     # set of obs to prepare for prediction (encode)
     clss_to_pred: Optional[list[str]] = field(default_factory=list)
     # set of obs that need to be hierarchically prepared
     hierarchical_clss: Optional[list[str]] = field(default_factory=list)
-    join_vars: Optional[Literal["auto", "inner", "None"]] = "None"
+    join_vars: Literal["inner", "outer"] | None = None
 
     def __post_init__(self):
-        self.mapped_dataset = mapped.mapped(
+        self.mapped_dataset = mapped(
             self.lamin_dataset,
-            label_keys=self.obs,
+            obs_keys=self.obs,
+            join=self.join_vars,
             encode_labels=self.clss_to_pred,
+            unknown_label="unknown",
             stream=True,
             parallel=True,
-            join_vars=self.join_vars,
         )
         print(
             "won't do any check but we recommend to have your dataset coming from local storage"
         )
-        self.class_groupings = {}
+        self.labels_groupings = {}
         self.class_topred = {}
         # generate tree from ontologies
         if len(self.hierarchical_clss) > 0:
@@ -93,24 +98,19 @@ class Dataset(torchDataset):
                     self.class_topred[clss] = self.mapped_dataset.get_merged_categories(
                         clss
                     )
-                    update = {}
-                    c = 0
-                    for k, v in self.mapped_dataset.encoders[clss].items():
-                        if k == self.mapped_dataset.unknown_class:
-                            update.update({k: v})
-                            c += 1
-                            self.class_topred[clss] -= set([k])
-                        else:
-                            update.update({k: v - c})
-                    self.mapped_dataset.encoders[clss] = update
+                    if (
+                        self.mapped_dataset.unknown_label
+                        in self.mapped_dataset.encoders[clss].keys()
+                    ):
+                        self.class_topred[clss] -= set(
+                            [self.mapped_dataset.unknown_label]
+                        )
 
         if self.genedf is None:
             self.genedf = load_genes(self.organisms)
 
         self.genedf.columns = self.genedf.columns.astype(str)
-        for organism in self.organisms:
-            ogenedf = self.genedf[self.genedf.organism == organism]
-            self.mapped_dataset._check_aligned_vars(ogenedf.index.tolist())
+        self.mapped_dataset._check_aligned_vars(self.genedf.index.tolist())
 
     def __len__(self, **kwargs):
         return self.mapped_dataset.__len__(**kwargs)
@@ -121,19 +121,6 @@ class Dataset(torchDataset):
 
     def __getitem__(self, *args, **kwargs):
         item = self.mapped_dataset.__getitem__(*args, **kwargs)
-        # import pdb
-
-        # pdb.set_trace()
-        # item.update(
-        #    {"unseen_genes": self.get_unseen_mapped_dataset_elements(*args, **kwargs)}
-        # )
-        # ret = {}
-        # ret["count"] = item[0]
-        # for i, val in enumerate(self.obs):
-        #    ret[val] = item[1][i]
-        ## mark unseen genes with a flag
-        ## send the associated
-        # print(item[0].shape)
         return item
 
     def __repr__(self):
@@ -148,7 +135,6 @@ class Dataset(torchDataset):
             + "     {} labels\n".format(len(self.obs))
             + "     {} clss_to_pred\n".format(len(self.clss_to_pred))
             + "     {} hierarchical_clss\n".format(len(self.hierarchical_clss))
-            + "     {} join_vars\n".format(len(self.join_vars))
             + "     {} organisms\n".format(len(self.organisms))
             + (
                 "dataset contains {} classes to predict\n".format(
@@ -160,17 +146,41 @@ class Dataset(torchDataset):
         )
 
     def get_label_weights(self, *args, **kwargs):
+        """
+        get_label_weights is a wrapper around mappedDataset.get_label_weights
+
+        Returns:
+            dict: dictionary of weights for each label
+        """
         return self.mapped_dataset.get_label_weights(*args, **kwargs)
 
     def get_unseen_mapped_dataset_elements(self, idx: int):
+        """
+        get_unseen_mapped_dataset_elements is a wrapper around mappedDataset.get_unseen_mapped_dataset_elements
+
+        Args:
+            idx (int): index of the element to get
+
+        Returns:
+            list[str]: list of unseen genes
+        """
         return [str(i)[2:-1] for i in self.mapped_dataset.uns(idx, "unseen_genes")]
 
-    def define_hierarchies(self, labels: list[str]):
+    def define_hierarchies(self, clsses: list[str]):
+        """
+        define_hierarchies is a method to define the hierarchies for the classes to predict
+
+        Args:
+            clsses (list[str]): list of classes to predict
+
+        Raises:
+            ValueError: if the class is not in the accepted classes
+        """
         # TODO: use all possible hierarchies instead of just the ones for which we have a sample annotated with
-        self.class_groupings = {}
+        self.labels_groupings = {}
         self.class_topred = {}
-        for label in labels:
-            if label not in [
+        for clss in clsses:
+            if clss not in [
                 "cell_type_ontology_term_id",
                 "tissue_ontology_term_id",
                 "disease_ontology_term_id",
@@ -179,41 +189,41 @@ class Dataset(torchDataset):
                 "self_reported_ethnicity_ontology_term_id",
             ]:
                 raise ValueError(
-                    "label {} not in accepted labels, for now only supported from bionty sources".format(
-                        label
+                    "class {} not in accepted classes, for now only supported from bionty sources".format(
+                        clss
                     )
                 )
-            elif label == "cell_type_ontology_term_id":
+            elif clss == "cell_type_ontology_term_id":
                 parentdf = (
                     bt.CellType.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
-            elif label == "tissue_ontology_term_id":
+            elif clss == "tissue_ontology_term_id":
                 parentdf = (
                     bt.Tissue.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
-            elif label == "disease_ontology_term_id":
+            elif clss == "disease_ontology_term_id":
                 parentdf = (
                     bt.Disease.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
-            elif label == "development_stage_ontology_term_id":
+            elif clss == "development_stage_ontology_term_id":
                 parentdf = (
                     bt.DevelopmentalStage.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
-            elif label == "assay_ontology_term_id":
+            elif clss == "assay_ontology_term_id":
                 parentdf = (
                     bt.ExperimentalFactor.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
-            elif label == "self_reported_ethnicity_ontology_term_id":
+            elif clss == "self_reported_ethnicity_ontology_term_id":
                 parentdf = (
                     bt.Ethnicity.filter()
                     .df(include=["parents__ontology_id"])
@@ -222,65 +232,58 @@ class Dataset(torchDataset):
 
             else:
                 raise ValueError(
-                    "label {} not in accepted labels, for now only supported from bionty sources".format(
-                        label
+                    "class {} not in accepted classes, for now only supported from bionty sources".format(
+                        clss
                     )
                 )
-            cats = self.mapped_dataset.get_merged_categories(label)
-            addition = set(LABELS_TOADD.get(label, {}).values())
+            cats = self.mapped_dataset.get_merged_categories(clss)
+            addition = set(LABELS_TOADD.get(clss, {}).values())
             cats |= addition
-            # import pdb
-
-            # pdb.set_trace()
-            groupings, _, lclass = get_ancestry_mapping(cats, parentdf)
+            groupings, _, leaf_labels = get_ancestry_mapping(cats, parentdf)
             for i, j in groupings.items():
                 if len(j) == 0:
                     groupings.pop(i)
-            self.class_groupings[label] = groupings
-            if label in self.clss_to_pred:
-                # if we have added new labels, we need to update the encoder with them too.
-                mlength = len(self.mapped_dataset.encoders[label])
+            self.labels_groupings[clss] = groupings
+            if clss in self.clss_to_pred:
+                # if we have added new clss, we need to update the encoder with them too.
+                mlength = len(self.mapped_dataset.encoders[clss])
+
                 mlength -= (
                     1
-                    if self.mapped_dataset.unknown_class
-                    in self.mapped_dataset.encoders[label].keys()
+                    if self.mapped_dataset.unknown_label
+                    in self.mapped_dataset.encoders[clss].keys()
                     else 0
                 )
 
                 for i, v in enumerate(
-                    addition - set(self.mapped_dataset.encoders[label].keys())
+                    addition - set(self.mapped_dataset.encoders[clss].keys())
                 ):
-                    self.mapped_dataset.encoders[label].update({v: mlength + i})
+                    self.mapped_dataset.encoders[clss].update({v: mlength + i})
                 # we need to change the ordering so that the things that can't be predicted appear afterward
 
-                self.class_topred[label] = lclass
+                self.class_topred[clss] = leaf_labels
                 c = 0
-                d = 0
                 update = {}
-                mlength = len(lclass)
-                # import pdb
-
-                # pdb.set_trace()
+                mlength = len(leaf_labels)
                 mlength -= (
                     1
-                    if self.mapped_dataset.unknown_class
-                    in self.mapped_dataset.encoders[label].keys()
+                    if self.mapped_dataset.unknown_label
+                    in self.mapped_dataset.encoders[clss].keys()
                     else 0
                 )
-                for k, v in self.mapped_dataset.encoders[label].items():
-                    if k in self.class_groupings[label].keys():
+                for k, v in self.mapped_dataset.encoders[clss].items():
+                    if k in self.labels_groupings[clss].keys():
                         update.update({k: mlength + c})
                         c += 1
-                    elif k == self.mapped_dataset.unknown_class:
+                    elif k == self.mapped_dataset.unknown_label:
                         update.update({k: v})
-                        d += 1
-                        self.class_topred[label] -= set([k])
+                        self.class_topred[clss] -= set([k])
                     else:
-                        update.update({k: (v - c) - d})
-                self.mapped_dataset.encoders[label] = update
+                        update.update({k: v - c})
+                self.mapped_dataset.encoders[clss] = update
 
 
-class SimpleAnnDataset:
+class SimpleAnnDataset(torchDataset):
     def __init__(
         self,
         adata: AnnData,
@@ -297,20 +300,65 @@ class SimpleAnnDataset:
             obs_to_output (list[str]): list of observations to output from anndata.obs
             layer (str): layer of the anndata to use
         """
-        self.adata = adata
-        self.obs_to_output = obs_to_output
-        self.layer = layer
+        self.adataX = adata.layers[layer] if layer is not None else adata.X
+        self.adataX = self.adataX.toarray() if issparse(self.adataX) else self.adataX
+        self.obs_to_output = adata.obs[obs_to_output]
 
     def __len__(self):
-        return self.adata.shape[0]
+        return self.adataX.shape[0]
+
+    def __iter__(self):
+        for idx, obs in enumerate(self.adata.obs.itertuples(index=False)):
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=DeprecationWarning)
+                out = {"X": self.adataX[idx].reshape(-1)}
+                out.update(
+                    {name: val for name, val in self.obs_to_output.iloc[idx].items()}
+                )
+                yield out
 
     def __getitem__(self, idx):
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=DeprecationWarning)
-            if self.layer is not None:
-                out = {"x": self.adata.layers[self.layer][idx].toarray().reshape(-1)}
-            else:
-                out = {"x": self.adata.X[idx].toarray().reshape(-1)}
-            for i in self.obs_to_output:
-                out.update({i: self.adata.obs.iloc[idx][i]})
+            out = {"X": self.adataX[idx].reshape(-1)}
+            out.update(
+                {name: val for name, val in self.obs_to_output.iloc[idx].items()}
+            )
         return out
+
+
+def mapped(
+    dataset,
+    obs_keys: list[str] | None = None,
+    join: Literal["inner", "outer"] | None = "inner",
+    encode_labels: bool | list[str] = True,
+    unknown_label: str | dict[str, str] | None = None,
+    cache_categories: bool = True,
+    parallel: bool = False,
+    dtype: str | None = None,
+    stream: bool = False,
+    is_run_input: bool | None = None,
+) -> MappedCollection:
+    path_list = []
+    for artifact in dataset.artifacts.all():
+        if artifact.suffix not in {".h5ad", ".zrad", ".zarr"}:
+            print(f"Ignoring artifact with suffix {artifact.suffix}")
+            continue
+        elif not artifact.path.exists():
+            print(f"Path does not exist for artifact with suffix {artifact.suffix}")
+            continue
+        elif not stream:
+            path_list.append(artifact.stage())
+        else:
+            path_list.append(artifact.path)
+    ds = MappedCollection(
+        path_list=path_list,
+        obs_keys=obs_keys,
+        join=join,
+        encode_labels=encode_labels,
+        unknown_label=unknown_label,
+        cache_categories=cache_categories,
+        parallel=parallel,
+        dtype=dtype,
+    )
+    return ds