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

scdataloader/data.py

@@ -1,83 +1,53 @@
 from dataclasses import dataclass, field
 
 import lamindb as ln
-import lnschema_bionty as lb
+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 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.
+
+    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, see :meth:`~lamindb.Dataset.mapped`.
+    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,16 +58,16 @@ 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: Optional[Literal["auto", "inner", "None"]] = "None"
 
     def __post_init__(self):
         self.mapped_dataset = mapped.mapped(
@@ -111,6 +81,8 @@ class Dataset(torchDataset):
         print(
             "won't do any check but we recommend to have your dataset coming from local storage"
         )
+        self.class_groupings = {}
+        self.class_topred = {}
         # generate tree from ontologies
         if len(self.hierarchical_clss) > 0:
             self.define_hierarchies(self.hierarchical_clss)
@@ -149,7 +121,12 @@ 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)})
+        # 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):
@@ -160,51 +137,36 @@ class Dataset(torchDataset):
         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))
+            + "     {} join_vars\n".format(len(self.join_vars))
+            + "     {} 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):
         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):
         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):
+    def define_hierarchies(self, labels: list[str]):
+        # TODO: use all possible hierarchies instead of just the ones for which we have a sample annotated with
         self.class_groupings = {}
         self.class_topred = {}
         for label in labels:
@@ -223,37 +185,37 @@ class Dataset(torchDataset):
                 )
             elif label == "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":
                 parentdf = (
-                    lb.Tissue.filter()
+                    bt.Tissue.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
             elif label == "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":
                 parentdf = (
-                    lb.DevelopmentalStage.filter()
+                    bt.DevelopmentalStage.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
             elif label == "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":
                 parentdf = (
-                    lb.Ethnicity.filter()
+                    bt.Ethnicity.filter()
                     .df(include=["parents__ontology_id"])
                     .set_index("ontology_id")
                 )
@@ -267,6 +229,9 @@ class Dataset(torchDataset):
             cats = self.mapped_dataset.get_merged_categories(label)
             addition = set(LABELS_TOADD.get(label, {}).values())
             cats |= addition
+            # import pdb
+
+            # pdb.set_trace()
             groupings, _, lclass = get_ancestry_mapping(cats, parentdf)
             for i, j in groupings.items():
                 if len(j) == 0:
@@ -316,7 +281,22 @@ class Dataset(torchDataset):
 
 
 class SimpleAnnDataset:
-    def __init__(self, adata, obs_to_output=[], layer=None):
+    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.adata = adata
         self.obs_to_output = obs_to_output
         self.layer = layer

scdataloader/datamodule.py

@@ -0,0 +1,375 @@
+import numpy as np
+import pandas as pd
+import lamindb as ln
+
+from torch.utils.data.sampler import (
+    WeightedRandomSampler,
+    SubsetRandomSampler,
+    SequentialSampler,
+)
+import torch
+from torch.utils.data import DataLoader, Sampler
+import lightning as L
+
+from typing import Optional, Union, Sequence
+
+from .data import Dataset
+from .collator import Collator
+from .utils import getBiomartTable
+
+
+class DataModule(L.LightningDataModule):
+    def __init__(
+        self,
+        collection_name: str,
+        label_to_weight: list = ["organism_ontology_term_id"],
+        organisms: list = ["NCBITaxon:9606"],
+        weight_scaler: int = 10,
+        train_oversampling_per_epoch: float = 0.1,
+        validation_split: float = 0.2,
+        test_split: float = 0,
+        gene_embeddings: str = "",
+        use_default_col: bool = True,
+        gene_position_tolerance: int = 10_000,
+        # this is for the mappedCollection
+        label_to_pred: list = ["organism_ontology_term_id"],
+        all_labels: list = ["organism_ontology_term_id"],
+        hierarchical_labels: list = [],
+        # this is for the collator
+        how: str = "random expr",
+        organism_name: str = "organism_ontology_term_id",
+        max_len: int = 1000,
+        add_zero_genes: int = 100,
+        do_gene_pos: Union[bool, str] = True,
+        tp_name: Optional[str] = None,  # "heat_diff"
+        assays_to_drop: list = [
+            "EFO:0008853",
+            "EFO:0010961",
+            "EFO:0030007",
+            "EFO:0030062",
+        ],
+        **kwargs,
+    ):
+        """
+        DataModule a pytorch lighting datamodule directly from a lamin Collection.
+        it can work with bare pytorch too
+
+        It implements train / val / test dataloaders. the train is weighted random, val is random, test is one to many separated datasets.
+        This is where the mappedCollection, dataset, and collator are combined to create the dataloaders.
+
+        Args:
+            collection_name (str): The lamindb collection to be used.
+            weight_scaler (int, optional): how much more you will see the most present vs less present category.
+            gene_position_tolerance (int, optional): The tolerance for gene position. Defaults to 10_000.
+                any genes within this distance of each other will be considered at the same position.
+            gene_embeddings (str, optional): The path to the gene embeddings file. Defaults to "".
+                the file must have ensembl_gene_id as index.
+                This is used to subset the available genes further to the ones that have embeddings in your model.
+            organisms (list, optional): The organisms to include in the dataset. Defaults to ["NCBITaxon:9606"].
+            label_to_weight (list, optional): List of labels to weight in the trainer's weighted random sampler. 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.
+                it will use a full dataset and will round to the nearest dataset's cell count.
+            **other args: see @file data.py and @file collator.py for more details
+            **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:
+            if type(do_gene_pos) is str:
+                print("seeing a string: loading gene positions as biomart parquet file")
+                biomart = pd.read_parquet(do_gene_pos)
+            else:
+                # 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.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=tp_name,
+                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.assays_to_drop = assays_to_drop
+        self.n_samples = len(mdataset)
+        self.weight_scaler = weight_scaler
+        self.train_oversampling_per_epoch = train_oversampling_per_epoch
+        self.label_to_weight = label_to_weight
+        self.train_weights = None
+        self.train_labels = None
+        super().__init__()
+
+    def __repr__(self):
+        return (
+            f"DataLoader(\n"
+            f"\twith a dataset=({self.dataset.__repr__()}\n)\n"
+            f"\tvalidation_split={self.validation_split},\n"
+            f"\ttest_split={self.test_split},\n"
+            f"\tn_samples={self.n_samples},\n"
+            f"\tweight_scaler={self.weight_scaler},\n"
+            f"\train_oversampling_per_epoch={self.train_oversampling_per_epoch},\n"
+            f"\tlabel_to_weight={self.label_to_weight}\n"
+            + (
+                "\twith train_dataset size of=("
+                + str((self.train_weights != 0).sum())
+                + ")\n)"
+            )
+            if self.train_weights is not None
+            else ")"
+        )
+
+    @property
+    def decoders(self):
+        """
+        decoders the decoders for any labels that would have been encoded
+
+        Returns:
+            dict[str, dict[int, str]]
+        """
+        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 the hierarchy of labels for any cls that would have a hierarchy
+
+        Returns:
+            dict[str, dict[str, str]]
+        """
+        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):
+        """
+        genes the genes used in this datamodule
+
+        Returns:
+            list
+        """
+        return self.dataset.genedf.index.tolist()
+
+    @property
+    def num_datasets(self):
+        return len(self.dataset.mapped_dataset.storages)
+
+    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, labels = self.dataset.get_label_weights(
+                self.label_to_weight, scaler=self.weight_scaler
+            )
+        else:
+            weights = np.ones(1)
+            labels = np.zeros(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 = []
+        if len(self.assays_to_drop) > 0:
+            for i, a in enumerate(
+                self.dataset.mapped_dataset.get_merged_labels("assay_ontology_term_id")
+            ):
+                if a not in self.assays_to_drop:
+                    idx_full.append(i)
+            idx_full = np.array(idx_full)
+        else:
+            idx_full = np.arange(self.n_samples)
+        test_datasets = []
+        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
+            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)
+            self.test_idx = idx_full[:len_test]
+            idx_full = idx_full[len_test:]
+        else:
+            self.test_idx = None
+
+        np.random.shuffle(idx_full)
+        if len_valid > 0:
+            self.valid_idx = idx_full[:len_valid].copy()
+            idx_full = idx_full[len_valid:]
+        else:
+            self.valid_idx = None
+        weights = np.concatenate([weights, np.zeros(1)])
+        labels[~np.isin(np.arange(self.n_samples), idx_full)] = len(weights) - 1
+
+        self.train_weights = weights
+        self.train_labels = labels
+        self.idx_full = idx_full
+
+        return test_datasets
+
+    def train_dataloader(self, **kwargs):
+        # train_sampler = WeightedRandomSampler(
+        #    self.train_weights[self.train_labels],
+        #    int(self.n_samples*self.train_oversampling_per_epoch),
+        #    replacement=True,
+        # )
+        train_sampler = LabelWeightedSampler(
+            self.train_weights,
+            self.train_labels,
+            num_samples=int(self.n_samples * self.train_oversampling_per_epoch),
+            # replacement=True,
+        )
+        return DataLoader(self.dataset, sampler=train_sampler, **self.kwargs, **kwargs)
+
+    def val_dataloader(self):
+        return (
+            DataLoader(
+                self.dataset, sampler=SubsetRandomSampler(self.valid_idx), **self.kwargs
+            )
+            if self.valid_idx is not None
+            else None
+        )
+
+    def test_dataloader(self):
+        return (
+            DataLoader(
+                self.dataset, sampler=SequentialSampler(self.test_idx), **self.kwargs
+            )
+            if self.test_idx is not None
+            else None
+        )
+
+    # def teardown(self):
+    # clean up state after the trainer stops, delete files...
+    # called on every process in DDP
+    # pass
+
+
+class LabelWeightedSampler(Sampler[int]):
+    label_weights: Sequence[float]
+    klass_indices: Sequence[Sequence[int]]
+    num_samples: int
+
+    # when we use, just set weights for each classes(here is: np.ones(num_classes)), and labels of a dataset.
+    # this will result a class-balanced sampling, no matter how imbalance the labels are.
+    # NOTE: here we use replacement=True, you can change it if you don't upsample a class.
+    def __init__(
+        self, label_weights: Sequence[float], labels: Sequence[int], num_samples: int
+    ) -> None:
+        """
+
+        :param label_weights: list(len=num_classes)[float], weights for each class.
+        :param labels: list(len=dataset_len)[int], labels of a dataset.
+        :param num_samples: number of samples.
+        """
+
+        super(LabelWeightedSampler, self).__init__(None)
+        # reweight labels from counter otherwsie same weight to labels that have many elements vs a few
+        label_weights = np.array(label_weights) * np.bincount(labels)
+
+        self.label_weights = torch.as_tensor(label_weights, dtype=torch.float32)
+        self.labels = torch.as_tensor(labels, dtype=torch.int)
+        self.num_samples = num_samples
+        # list of tensor.
+        self.klass_indices = [
+            (self.labels == i_klass).nonzero().squeeze(1)
+            for i_klass in range(len(label_weights))
+        ]
+
+    def __iter__(self):
+        sample_labels = torch.multinomial(
+            self.label_weights, num_samples=self.num_samples, replacement=True
+        )
+        sample_indices = torch.empty_like(sample_labels)
+        for i_klass, klass_index in enumerate(self.klass_indices):
+            if klass_index.numel() == 0:
+                continue
+            left_inds = (sample_labels == i_klass).nonzero().squeeze(1)
+            right_inds = torch.randint(len(klass_index), size=(len(left_inds),))
+            sample_indices[left_inds] = klass_index[right_inds]
+        yield from iter(sample_indices.tolist())
+
+    def __len__(self):
+        return self.num_samples