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

scdataloader/data.py

@@ -4,46 +4,47 @@ import lamindb as ln
 import lnschema_bionty as lb
 import pandas as pd
 from torch.utils.data import Dataset as torchDataset
-
+from typing import Union
 from scdataloader import mapped
+import warnings
 
 # TODO: manage load gene embeddings to make
 # from scprint.dataloader.embedder import embed
-from scdataloader.utils import get_ancestry_mapping, pd_load_cached
+from scdataloader.utils import get_ancestry_mapping, load_genes
 
 LABELS_TOADD = {
-    "assay_ontology_term_id": [
-        "10x transcription profiling",
-        "spatial transcriptomics",
-        "10x 3' transcription profiling",
-        "10x 5' transcription profiling",
-    ],
-    "disease_ontology_term_id": [
-        "metabolic disease",
-        "chronic kidney disease",
-        "chromosomal disorder",
-        "infectious disease",
-        "inflammatory disease",
+    "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",
-        "mitochondrial disease",
-        "psychiatric disorder",
-        "cancer or benign tumor",
-        "neoplasm",
-    ],
-    "cell_type_ontology_term_id": [
-        "progenitor cell",
-        "hematopoietic cell",
-        "myoblast",
-        "myeloid cell",
-        "neuron",
-        "electrically active cell",
-        "epithelial cell",
-        "secretory cell",
-        "stem cell",
-        "non-terminally differentiated cell",
-        "supporting cell",
-    ],
+        "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",
+    },
 }
 
 
@@ -66,16 +67,14 @@ class Dataset(torchDataset):
         gene_embedding: dataframe containing the gene embeddings
         organisms (list[str]): list of organisms to load
         obs (list[str]): list of observations to load
-        encode_obs (list[str]): list of observations to encode
-        map_hierarchy: list of observations to map to a hierarchy
+        clss_to_pred (list[str]): list of observations to encode
+        hierarchical_clss: list of observations to map to a hierarchy
     """
 
-    lamin_dataset: ln.Dataset
+    lamin_dataset: ln.Collection
     genedf: pd.DataFrame = None
-    gene_embedding: pd.DataFrame = (
-        None  # TODO: make it part of specialized dataset
-    )
-    organisms: list[str] = field(
+    # gene_embedding: pd.DataFrame = None  # TODO: make it part of specialized dataset
+    organisms: Union[list[str], str] = field(
         default_factory=["NCBITaxon:9606", "NCBITaxon:10090"]
     )
     obs: list[str] = field(
@@ -94,43 +93,63 @@ class Dataset(torchDataset):
             "nnz",
         ]
     )
-    encode_obs: list[str] = field(default_factory=list)
-    map_hierarchy: list[str] = field(default_factory=list)
+    # set of obs to prepare for prediction (encode)
+    clss_to_pred: 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"
 
     def __post_init__(self):
         self.mapped_dataset = mapped.mapped(
             self.lamin_dataset,
             label_keys=self.obs,
-            encode_labels=self.encode_obs,
+            encode_labels=self.clss_to_pred,
             stream=True,
             parallel=True,
-            join_vars="None",
+            join_vars=self.join_vars,
         )
         print(
             "won't do any check but we recommend to have your dataset coming from local storage"
         )
         # generate tree from ontologies
-        if len(self.map_hierarchy) > 0:
-            self.define_hierarchies(self.map_hierarchy)
+        if len(self.hierarchical_clss) > 0:
+            self.define_hierarchies(self.hierarchical_clss)
+        if len(self.clss_to_pred) > 0:
+            for clss in self.clss_to_pred:
+                if clss not in self.hierarchical_clss:
+                    # otherwise it's already been done
+                    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.genedf is None:
-            self.genedf = self.load_genes(self.organisms)
+            self.genedf = load_genes(self.organisms)
 
-        if self.gene_embedding is None:
-            self.gene_embedding = self.load_embeddings(self.genedf)
-        else:
-            # self.genedf = pd.concat(
-            #    [self.genedf.set_index("ensembl_gene_id"), self.gene_embedding],
-            #    axis=1,
-            #    join="inner",
-            # )
-            self.genedf.columns = self.genedf.columns.astype(str)
+        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())
 
     def __len__(self, **kwargs):
         return self.mapped_dataset.__len__(**kwargs)
 
+    @property
+    def encoder(self):
+        return self.mapped_dataset.encoders
+
     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):
@@ -143,8 +162,7 @@ class Dataset(torchDataset):
     def __repr__(self):
         print(
             "total dataset size is {} Gb".format(
-                sum([file.size for file in self.lamin_dataset.artifacts.all()])
-                / 1e9
+                sum([file.size for file in self.lamin_dataset.artifacts.all()]) / 1e9
             )
         )
         print("---")
@@ -158,111 +176,14 @@ class Dataset(torchDataset):
                 sum([len(self.class_topred[i]) for i in self.class_topred])
             )
         )
-        print("embedding size is {}".format(self.gene_embedding.shape[1]))
+        # print("embedding size is {}".format(self.gene_embedding.shape[1]))
         return ""
 
     def get_label_weights(self, *args, **kwargs):
         return self.mapped_dataset.get_label_weights(*args, **kwargs)
 
     def get_unseen_mapped_dataset_elements(self, idx):
-        return [
-            str(i)[2:-1] for i in self.mapped_dataset.uns(idx, "unseen_genes")
-        ]
-
-    def use_prior_network(
-        self, name="collectri", organism="human", split_complexes=True
-    ):
-        """
-        use_prior_network loads a prior GRN from a list of available networks.
-
-        Args:
-            name (str, optional): name of the network to load. Defaults to "collectri".
-            organism (str, optional): organism to load the network for. Defaults to "human".
-            split_complexes (bool, optional): whether to split complexes into individual genes. Defaults to True.
-
-        Raises:
-            ValueError: if the provided name is not amongst the available names.
-        """
-        # TODO: use omnipath instead
-        if name == "tflink":
-            TFLINK = "https://cdn.netbiol.org/tflink/download_files/TFLink_Homo_sapiens_interactions_All_simpleFormat_v1.0.tsv.gz"
-            net = pd_load_cached(TFLINK)
-            net = net.rename(
-                columns={"Name.TF": "regulator", "Name.Target": "target"}
-            )
-        elif name == "htftarget":
-            HTFTARGET = "http://bioinfo.life.hust.edu.cn/static/hTFtarget/file_download/tf-target-infomation.txt"
-            net = pd_load_cached(HTFTARGET)
-            net = net.rename(columns={"TF": "regulator"})
-        elif name == "collectri":
-            import decoupler as dc
-
-            net = dc.get_collectri(
-                organism=organism, split_complexes=split_complexes
-            )
-            net = net.rename(columns={"source": "regulator"})
-        else:
-            raise ValueError(
-                f"provided name: '{name}' is not amongst the available names."
-            )
-        self.add_prior_network(net)
-
-    def add_prior_network(self, prior_network: pd.DataFrame, init_len):
-        # validate the network dataframe
-        required_columns: list[str] = ["target", "regulators"]
-        optional_columns: list[str] = ["type", "weight"]
-
-        for column in required_columns:
-            assert (
-                column in prior_network.columns
-            ), f"Column '{column}' is missing in the provided network dataframe."
-
-        for column in optional_columns:
-            if column not in prior_network.columns:
-                print(
-                    f"Optional column '{column}' is not present in the provided network dataframe."
-                )
-
-        assert (
-            prior_network["target"].dtype == "str"
-        ), "Column 'target' should be of dtype 'str'."
-        assert (
-            prior_network["regulators"].dtype == "str"
-        ), "Column 'regulators' should be of dtype 'str'."
-
-        if "type" in prior_network.columns:
-            assert (
-                prior_network["type"].dtype == "str"
-            ), "Column 'type' should be of dtype 'str'."
-
-        if "weight" in prior_network.columns:
-            assert (
-                prior_network["weight"].dtype == "float"
-            ), "Column 'weight' should be of dtype 'float'."
-
-        # TODO: check that we match the genes in the network to the genes in the dataset
-
-        print(
-            "loaded {:.2f}% of the edges".format(
-                (len(prior_network) / init_len) * 100
-            )
-        )
-        # TODO: transform it into a sparse matrix
-        self.prior_network = prior_network
-        self.network_size = len(prior_network)
-        # self.overlap =
-        # self.edge_freq
-
-    def load_genes(self, organisms):
-        organismdf = []
-        for o in organisms:
-            organism = lb.Gene(
-                organism=lb.Organism.filter(ontology_id=o).one()
-            ).df()
-            organism["organism"] = o
-            organismdf.append(organism)
-
-        return pd.concat(organismdf)
+        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 = []
@@ -344,10 +265,72 @@ class Dataset(torchDataset):
                     )
                 )
             cats = self.mapped_dataset.get_merged_categories(label)
-            cats |= set(LABELS_TOADD.get(label, []))
+            addition = set(LABELS_TOADD.get(label, {}).values())
+            cats |= addition
             groupings, _, lclass = get_ancestry_mapping(cats, parentdf)
             for i, j in groupings.items():
                 if len(j) == 0:
                     groupings.pop(i)
             self.class_groupings[label] = groupings
-            self.class_topred[label] = lclass
+            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])
+                mlength -= (
+                    1
+                    if self.mapped_dataset.unknown_class
+                    in self.mapped_dataset.encoders[label].keys()
+                    else 0
+                )
+
+                for i, v in enumerate(
+                    addition - set(self.mapped_dataset.encoders[label].keys())
+                ):
+                    self.mapped_dataset.encoders[label].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
+                c = 0
+                d = 0
+                update = {}
+                mlength = len(lclass)
+                # import pdb
+
+                # pdb.set_trace()
+                mlength -= (
+                    1
+                    if self.mapped_dataset.unknown_class
+                    in self.mapped_dataset.encoders[label].keys()
+                    else 0
+                )
+                for k, v in self.mapped_dataset.encoders[label].items():
+                    if k in self.class_groupings[label].keys():
+                        update.update({k: mlength + c})
+                        c += 1
+                    elif k == self.mapped_dataset.unknown_class:
+                        update.update({k: v})
+                        d += 1
+                        self.class_topred[label] -= set([k])
+                    else:
+                        update.update({k: (v - c) - d})
+                self.mapped_dataset.encoders[label] = 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
+
+    def __len__(self):
+        return self.adata.shape[0]
+
+    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]})
+        return out