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

scdataloader/dataloader.py

@@ -0,0 +1,318 @@
+import numpy as np
+import pandas as pd
+import lamindb as ln
+
+from torch.utils.data.sampler import (
+    WeightedRandomSampler,
+    SubsetRandomSampler,
+    SequentialSampler,
+)
+from torch.utils.data import DataLoader
+import lightning as L
+
+from typing import Optional
+
+from .data import Dataset
+from .collator import Collator
+from .mapped import MappedDataset
+from .utils import getBiomartTable
+
+# TODO: put in config
+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": "",
+}
+
+
+class DataModule(L.LightningDataModule):
+    """
+    Base class for all data loaders
+    """
+
+    def __init__(
+        self,
+        mdataset: Optional[MappedDataset] = None,
+        collection_name=None,
+        organisms: list = ["NCBITaxon:9606"],
+        weight_scaler: int = 30,
+        train_oversampling=1,
+        label_to_weight: list = [],
+        label_to_pred: list = [],
+        validation_split: float = 0.2,
+        test_split: float = 0,
+        use_default_col=True,
+        all_labels=[],
+        hierarchical_labels=[],
+        how="most expr",
+        organism_name="organism_ontology_term_id",
+        max_len=1000,
+        add_zero_genes=100,
+        do_gene_pos=True,
+        gene_embeddings="",
+        gene_position_tolerance=10_000,
+        **kwargs,
+    ):
+        """
+        Initializes the DataModule.
+
+        Args:
+            dataset (MappedDataset): The dataset to be used.
+            weight_scaler (int, optional): The weight scaler for weighted random sampling. Defaults to 30.
+            label_to_weight (list, optional): List of labels to weight. Defaults to [].
+            validation_split (float, optional): The proportion of the dataset to include in the validation split. Defaults to 0.2.
+            test_split (float, optional): The proportion of the dataset to include in the test split. Defaults to 0.
+            **kwargs: Additional keyword arguments passed to the pytorch DataLoader.
+        """
+        if collection_name is not None:
+            mdataset = Dataset(
+                ln.Collection.filter(name=collection_name).first(),
+                organisms=organisms,
+                obs=all_labels,
+                clss_to_pred=label_to_pred,
+                hierarchical_clss=hierarchical_labels,
+            )
+            print(mdataset)
+        # and location
+        if do_gene_pos:
+            # and annotations
+            biomart = getBiomartTable(
+                attributes=["start_position", "chromosome_name"]
+            ).set_index("ensembl_gene_id")
+            biomart = biomart.loc[~biomart.index.duplicated(keep="first")]
+            biomart = biomart.sort_values(by=["chromosome_name", "start_position"])
+            c = []
+            i = 0
+            prev_position = -100000
+            prev_chromosome = None
+            for _, r in biomart.iterrows():
+                if (
+                    r["chromosome_name"] != prev_chromosome
+                    or r["start_position"] - prev_position > gene_position_tolerance
+                ):
+                    i += 1
+                c.append(i)
+                prev_position = r["start_position"]
+                prev_chromosome = r["chromosome_name"]
+            print(f"reduced the size to {len(set(c))/len(biomart)}")
+            biomart["pos"] = c
+            mdataset.genedf = biomart.loc[
+                mdataset.genedf[mdataset.genedf.index.isin(biomart.index)].index
+            ]
+            self.gene_pos = mdataset.genedf["pos"].tolist()
+
+        if gene_embeddings != "":
+            mdataset.genedf = mdataset.genedf.join(
+                pd.read_parquet(gene_embeddings), how="inner"
+            )
+            if do_gene_pos:
+                self.gene_pos = mdataset.genedf["pos"].tolist()
+        self.labels = {k: len(v) for k, v in mdataset.class_topred.items()}
+        # we might want not to order the genes by expression (or do it?)
+        # we might want to not introduce zeros and
+        if use_default_col:
+            kwargs["collate_fn"] = Collator(
+                organisms=organisms,
+                how=how,
+                valid_genes=mdataset.genedf.index.tolist(),
+                max_len=max_len,
+                add_zero_genes=add_zero_genes,
+                org_to_id=mdataset.encoder[organism_name],
+                tp_name="heat_diff",
+                organism_name=organism_name,
+                class_names=label_to_weight,
+            )
+        self.validation_split = validation_split
+        self.test_split = test_split
+        self.dataset = mdataset
+        self.kwargs = kwargs
+        self.n_samples = len(mdataset)
+        self.weight_scaler = weight_scaler
+        self.train_oversampling = train_oversampling
+        self.label_to_weight = label_to_weight
+        super().__init__()
+
+    @property
+    def decoders(self):
+        decoders = {}
+        for k, v in self.dataset.encoder.items():
+            decoders[k] = {va: ke for ke, va in v.items()}
+        return decoders
+
+    @property
+    def cls_hierarchy(self):
+        cls_hierarchy = {}
+        for k, dic in self.dataset.class_groupings.items():
+            rdic = {}
+            for sk, v in dic.items():
+                rdic[self.dataset.encoder[k][sk]] = [
+                    self.dataset.encoder[k][i] for i in list(v)
+                ]
+            cls_hierarchy[k] = rdic
+        return cls_hierarchy
+
+    @property
+    def genes(self):
+        return self.dataset.genedf.index.tolist()
+
+    def setup(self, stage=None):
+        """
+        setup method is used to prepare the data for the training, validation, and test sets.
+        It shuffles the data, calculates weights for each set, and creates samplers for each set.
+
+        Args:
+            stage (str, optional): The stage of the model training process.
+            It can be either 'fit' or 'test'. Defaults to None.
+        """
+
+        if len(self.label_to_weight) > 0:
+            weights = self.dataset.get_label_weights(
+                self.label_to_weight, scaler=self.weight_scaler
+            )
+        else:
+            weights = np.ones(self.n_samples)
+        if isinstance(self.validation_split, int):
+            len_valid = self.validation_split
+        else:
+            len_valid = int(self.n_samples * self.validation_split)
+        if isinstance(self.test_split, int):
+            len_test = self.test_split
+        else:
+            len_test = int(self.n_samples * self.test_split)
+        assert (
+            len_test + len_valid < self.n_samples
+        ), "test set + valid set size is configured to be larger than entire dataset."
+        idx_full = np.arange(self.n_samples)
+        if len_test > 0:
+            # this way we work on some never seen datasets
+            # keeping at least one
+            len_test = (
+                len_test
+                if len_test > self.dataset.mapped_dataset.n_obs_list[0]
+                else self.dataset.mapped_dataset.n_obs_list[0]
+            )
+            cs = 0
+            test_datasets = []
+            print("these files will be considered test datasets:")
+            for i, c in enumerate(self.dataset.mapped_dataset.n_obs_list):
+                if cs + c > len_test:
+                    break
+                else:
+                    print("    " + self.dataset.mapped_dataset.path_list[i].path)
+                    test_datasets.append(self.dataset.mapped_dataset.path_list[i].path)
+                    cs += c
+
+            len_test = cs
+            print("perc test: ", len_test / self.n_samples)
+            test_idx = idx_full[:len_test]
+            idx_full = idx_full[len_test:]
+            self.test_sampler = SequentialSampler(test_idx)
+        else:
+            self.test_sampler = None
+            test_datasets = None
+
+        np.random.shuffle(idx_full)
+        if len_valid > 0:
+            valid_idx = idx_full[:len_valid]
+            idx_full = idx_full[len_valid:]
+            self.valid_sampler = SubsetRandomSampler(valid_idx)
+        else:
+            self.valid_sampler = None
+
+        weights[~idx_full] = 0
+        self.train_sampler = WeightedRandomSampler(
+            weights,
+            int(len(idx_full) * self.train_oversampling),
+            replacement=True,
+        )
+        return test_datasets
+
+    def train_dataloader(self, **kwargs):
+        return DataLoader(
+            self.dataset, sampler=self.train_sampler, **self.kwargs, **kwargs
+        )
+
+    def val_dataloader(self):
+        return (
+            DataLoader(self.dataset, sampler=self.valid_sampler, **self.kwargs)
+            if self.valid_sampler is not None
+            else None
+        )
+
+    def test_dataloader(self):
+        return (
+            DataLoader(self.dataset, sampler=self.test_sampler, **self.kwargs)
+            if self.test_sampler is not None
+            else None
+        )
+
+    # def teardown(self):
+    # clean up state after the trainer stops, delete files...
+    # called on every process in DDP
+    # pass

scdataloader/mapped.py

@@ -80,10 +80,13 @@ class MappedDataset:
         join_vars: Optional[Literal["auto", "inner", "None"]] = "auto",
         encode_labels: Optional[Union[bool, List[str]]] = False,
         parallel: bool = False,
+        unknown_class: str = "unknown",
     ):
         self.storages = []
         self.conns = []
         self.parallel = parallel
+        self.unknown_class = unknown_class
+        self.path_list = path_list
         self._make_connections(path_list, parallel)
 
         self.n_obs_list = []
@@ -96,22 +99,16 @@ class MappedDataset:
                     self.n_obs_list.append(X.attrs["shape"][0])
         self.n_obs = sum(self.n_obs_list)
 
-        self.indices = np.hstack(
-            [np.arange(n_obs) for n_obs in self.n_obs_list]
-        )
-        self.storage_idx = np.repeat(
-            np.arange(len(self.storages)), self.n_obs_list
-        )
+        self.indices = np.hstack([np.arange(n_obs) for n_obs in self.n_obs_list])
+        self.storage_idx = np.repeat(np.arange(len(self.storages)), self.n_obs_list)
 
         self.join_vars = join_vars if len(path_list) > 1 else None
         self.var_indices = None
-        if self.join_vars is not None:
+        if self.join_vars != "None":
             self._make_join_vars()
 
         self.encode_labels = encode_labels
-        self.label_keys = (
-            [label_keys] if isinstance(label_keys, str) else label_keys
-        )
+        self.label_keys = [label_keys] if isinstance(label_keys, str) else label_keys
         if isinstance(encode_labels, bool):
             if encode_labels:
                 encode_labels = label_keys
@@ -122,6 +119,8 @@ class MappedDataset:
             for label in encode_labels:
                 cats = self.get_merged_categories(label)
                 self.encoders[label] = {cat: i for i, cat in enumerate(cats)}
+                if unknown_class in self.encoders[label]:
+                    self.encoders[label][unknown_class] = -1
         else:
             self.encoders = {}
         self._closed = False
@@ -157,9 +156,15 @@ class MappedDataset:
                 raise ValueError(
                     "The provided AnnData objects don't have shared varibales."
                 )
-            self.var_indices = [
-                vrs.get_indexer(self.var_joint) for vrs in var_list
-            ]
+            self.var_indices = [vrs.get_indexer(self.var_joint) for vrs in var_list]
+
+    def _check_aligned_vars(self, vars: list):
+        i = 0
+        for storage in self.storages:
+            with _Connect(storage) as store:
+                if vars == _safer_read_index(store["var"]).tolist():
+                    i += 1
+        print("{}% are aligned".format(i * 100 / len(self.storages)))
 
     def __len__(self):
         return self.n_obs
@@ -172,14 +177,14 @@ class MappedDataset:
         else:
             var_idxs = None
         with _Connect(self.storages[storage_idx]) as store:
-            out = [self.get_data_idx(store, obs_idx, var_idxs)]
+            out = {"x": self.get_data_idx(store, obs_idx, var_idxs)}
             if self.label_keys is not None:
                 for i, label in enumerate(self.label_keys):
                     label_idx = self.get_label_idx(store, obs_idx, label)
                     if label in self.encoders:
-                        out.append(self.encoders[label][label_idx])
+                        out.update({label: self.encoders[label][label_idx]})
                     else:
-                        out.append(label_idx)
+                        out.update({label: label_idx})
         return out
 
     def uns(self, idx, key):
@@ -240,9 +245,7 @@ class MappedDataset:
             if i == 0:
                 labels = self.get_merged_labels(val)
             else:
-                labels += "_" + self.get_merged_labels(val).astype(str).astype(
-                    "O"
-                )
+                labels += "_" + self.get_merged_labels(val).astype(str).astype("O")
         counter = Counter(labels)  # type: ignore
         counter = np.array([counter[label] for label in labels])
         weights = scaler / (counter + scaler)
@@ -255,9 +258,7 @@ class MappedDataset:
         for storage in self.storages:
             with _Connect(storage) as store:
                 codes = self.get_codes(store, label_key)
-                labels = (
-                    decode(codes) if isinstance(codes[0], bytes) else codes
-                )
+                labels = decode(codes) if isinstance(codes[0], bytes) else codes
                 cats = self.get_categories(store, label_key)
                 if cats is not None:
                     cats = decode(cats) if isinstance(cats[0], bytes) else cats
@@ -277,9 +278,7 @@ class MappedDataset:
                     cats_merge.update(cats)
                 else:
                     codes = self.get_codes(store, label_key)
-                    codes = (
-                        decode(codes) if isinstance(codes[0], bytes) else codes
-                    )
+                    codes = decode(codes) if isinstance(codes[0], bytes) else codes
                     cats_merge.update(codes)
         return cats_merge