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

scdataloader/data.py

@@ -1,83 +1,57 @@
 from dataclasses import dataclass, field
 
 import lamindb as ln
-import lnschema_bionty as lb
+
+# ln.connect("scprint")
+
+import bionty as bt
 import pandas as pd
 from torch.utils.data import Dataset as torchDataset
-from typing import Union
+from typing import Union, Optional, Literal
 from scdataloader import mapped
 import warnings
 
-# TODO: manage load gene embeddings to make
-# from scprint.dataloader.embedder import embed
+from anndata import AnnData
+from scipy.sparse import issparse
+
 from scdataloader.utils import get_ancestry_mapping, load_genes
 
-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",
-    },
-}
+from .config import LABELS_TOADD
 
 
 @dataclass
 class Dataset(torchDataset):
     """
-    Dataset class to load a bunch of anndata from a lamin dataset in a memory efficient way.
+    Dataset class to load a bunch of anndata from a lamin dataset (Collection) in a memory efficient way.
 
-    For an example, see :meth:`~lamindb.Dataset.mapped`.
+    This serves as a wrapper around lamin's mappedCollection to provide more features,
+    mostly, the management of hierarchical labels, the encoding of labels, the management of multiple species
+
+    For an example of mappedDataset, see :meth:`~lamindb.Dataset.mapped`.
 
     .. note::
 
-        A similar data loader exists `here
+        A related data loader exists `here
         <https://github.com/Genentech/scimilarity>`__.
 
-    Attributes:
+    Args:
     ----
         lamin_dataset (lamindb.Dataset): lamin dataset to load
         genedf (pd.Dataframe): dataframe containing the genes to load
-        gene_embedding: dataframe containing the gene embeddings
         organisms (list[str]): list of organisms to load
-        obs (list[str]): list of observations to load
+            (for now only validates the the genes map to this organism)
+        obs (list[str]): list of observations to load from the Collection
         clss_to_pred (list[str]): list of observations to encode
-        hierarchical_clss: list of observations to map to a hierarchy
+        join_vars (flag): join variables @see :meth:`~lamindb.Dataset.mapped`.
+        hierarchical_clss: list of observations to map to a hierarchy using lamin's bionty
     """
 
     lamin_dataset: ln.Collection
-    genedf: pd.DataFrame = None
-    # gene_embedding: pd.DataFrame = None  # TODO: make it part of specialized dataset
-    organisms: Union[list[str], str] = field(
+    genedf: Optional[pd.DataFrame] = None
+    organisms: Optional[Union[list[str], str]] = field(
         default_factory=["NCBITaxon:9606", "NCBITaxon:10090"]
     )
-    obs: list[str] = field(
+    obs: Optional[list[str]] = field(
         default_factory=[
             "self_reported_ethnicity_ontology_term_id",
             "assay_ontology_term_id",
@@ -88,29 +62,32 @@ 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: list[str] = field(default_factory=list)
+    clss_to_pred: Optional[list[str]] = field(default_factory=list)
     # set of obs that need to be hierarchically prepared
-    hierarchical_clss: list[str] = field(default_factory=list)
-    join_vars: str = "None"
+    hierarchical_clss: Optional[list[str]] = field(default_factory=list)
+    join_vars: Literal["inner", "outer"] | None = None
 
     def __post_init__(self):
         self.mapped_dataset = mapped.mapped(
             self.lamin_dataset,
             label_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.labels_groupings = {}
+        self.class_topred = {}
         # generate tree from ontologies
         if len(self.hierarchical_clss) > 0:
             self.define_hierarchies(self.hierarchical_clss)
@@ -121,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)
@@ -149,66 +121,66 @@ class Dataset(torchDataset):
 
     def __getitem__(self, *args, **kwargs):
         item = self.mapped_dataset.__getitem__(*args, **kwargs)
-        #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):
-        print(
-            "total dataset size is {} Gb".format(
+        return (
+            "total dataset size is {} Gb\n".format(
                 sum([file.size for file in self.lamin_dataset.artifacts.all()]) / 1e9
             )
-        )
-        print("---")
-        print("dataset contains:")
-        print("     {} cells".format(self.mapped_dataset.__len__()))
-        print("     {} genes".format(self.genedf.shape[0]))
-        print("     {} labels".format(len(self.obs)))
-        print("     {} organisms".format(len(self.organisms)))
-        print(
-            "dataset contains {} classes to predict".format(
-                sum([len(self.class_topred[i]) for i in self.class_topred])
+            + "---\n"
+            + "dataset contains:\n"
+            + "     {} cells\n".format(self.mapped_dataset.__len__())
+            + "     {} genes\n".format(self.genedf.shape[0])
+            + "     {} labels\n".format(len(self.obs))
+            + "     {} clss_to_pred\n".format(len(self.clss_to_pred))
+            + "     {} hierarchical_clss\n".format(len(self.hierarchical_clss))
+            + "     {} organisms\n".format(len(self.organisms))
+            + (
+                "dataset contains {} classes to predict\n".format(
+                    sum([len(self.class_topred[i]) for i in self.class_topred])
+                )
+                if len(self.class_topred) > 0
+                else ""
             )
         )
-        # print("embedding size is {}".format(self.gene_embedding.shape[1]))
-        return ""
 
     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):
+    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 load_embeddings(self, genedfs, embedding_size=128, cache=True):
-    #    embeddings = []
-    #    for o in self.organisms:
-    #        genedf = genedfs[genedfs.organism == o]
-    #        org_name = lb.Organism.filter(ontology_id=o).one().scientific_name
-    #        embedding = embed(
-    #            genedf=genedf,
-    #            organism=org_name,
-    #            cache=cache,
-    #            fasta_path="/tmp/data/fasta/",
-    #            embedding_size=embedding_size,
-    #        )
-    #        genedf = pd.concat(
-    #            [genedf.set_index("ensembl_gene_id"), embedding], axis=1, join="inner"
-    #        )
-    #        genedf.columns = genedf.columns.astype(str)
-    #        embeddings.append(genedf)
-    #    return pd.concat(embeddings)
-
-    def define_hierarchies(self, labels):
-        self.class_groupings = {}
+    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.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",
@@ -217,120 +189,139 @@ 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 = (
-                    lb.CellType.filter()
+                    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 = (
-                    lb.Tissue.filter()
+                    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 = (
-                    lb.Disease.filter()
+                    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 = (
-                    lb.DevelopmentalStage.filter()
+                    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 = (
-                    lb.ExperimentalFactor.filter()
+                    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 = (
-                    lb.Ethnicity.filter()
+                    bt.Ethnicity.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
 
             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
-            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:
-    def __init__(self, adata, obs_to_output=[], layer=None):
-        self.adata = adata
-        self.obs_to_output = obs_to_output
-        self.layer = layer
+class SimpleAnnDataset(torchDataset):
+    def __init__(
+        self,
+        adata: AnnData,
+        obs_to_output: Optional[list[str]] = [],
+        layer: Optional[str] = None,
+    ):
+        """
+        SimpleAnnDataset is a simple dataloader for an AnnData dataset. this is to interface nicely with the rest of
+        scDataloader and with your model during inference.
+
+        Args:
+        ----
+            adata (anndata.AnnData): anndata object to use
+            obs_to_output (list[str]): list of observations to output from anndata.obs
+            layer (str): layer of the anndata to use
+        """
+        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