Documentation for DataModule

def predict_dataloader(self):
    """
    Create a DataLoader for prediction over all training data.

    Uses RankShardSampler for distributed inference.

    Returns:
        DataLoader: Prediction DataLoader instance.
    """
    subset = Subset(self.dataset, self.idx_full)
    return DataLoader(
        self.dataset,
        sampler=RankShardSampler(len(subset), start_at=self.start_at),
        **self.kwargs,
    )

def setup(self, stage=None):
    """
    Prepare data splits for training, validation, and testing.

    This method shuffles the data, computes sample weights for weighted sampling,
    removes samples from dropped assays, and creates train/val/test splits.
    Test splits use entire datasets to ensure evaluation on unseen data sources.

    Results can be cached to `store_location` for faster subsequent runs.

    Args:
        stage (str, optional): Training stage ('fit', 'test', or None for both).
            Currently not used but kept for Lightning compatibility. Defaults to None.

    Returns:
        List[str]: List of paths to test datasets.

    Note:
        Must be called before using dataloaders. The train/val/test split is
        deterministic when loading from cache.
    """
    print("setting up the datamodule")
    start_time = time.time()
    if (
        self.store_location is None
        or not os.path.exists(os.path.join(self.store_location, "train_labels.npy"))
        or self.force_recompute_indices
    ):
        if "nnz" in self.clss_to_weight and self.weight_scaler > 0:
            self.nnz = self.dataset.mapped_dataset.get_merged_labels(
                "nnz", is_cat=False
            )
            self.clss_to_weight.remove("nnz")
            # Sigmoid scaling with 2 parameters
            midpoint = 2000
            steepness = 0.003
            # Apply sigmoid transformation
            # sigmoid(x) = 1 / (1 + exp(-steepness * (x - midpoint)))
            # Then scale to [1, NNZ_SCALE] range
            sigmoid_values = 1 / (1 + np.exp(-steepness * (self.nnz - midpoint)))
            self.nnz = 1 + ((NNZ_SCALE - 1) * sigmoid_values)
        if len(self.clss_to_weight) > 0 and self.weight_scaler > 0:
            labels = self.dataset.get_label_cats(
                self.clss_to_weight,
            )
        else:
            labels = np.zeros(self.n_samples, dtype=int)
        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:
            badloc = np.isin(
                self.dataset.mapped_dataset.get_merged_labels(
                    "assay_ontology_term_id"
                ),
                self.assays_to_drop,
            )
            idx_full = np.arange(len(labels))[~badloc]
        else:
            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
            d_size = list(enumerate(self.dataset.mapped_dataset.n_obs_list))
            random.Random(42).shuffle(d_size)  # always same order
            for i, c in d_size:
                if cs + c > len_test:
                    break
                else:
                    self.test_datasets.append(
                        self.dataset.mapped_dataset.path_list[i].path
                    )
                    cs += c
            len_test = cs
            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()
            # store it for later
            idx_full = idx_full[len_valid:]
        else:
            self.valid_idx = None
        labels[~np.isin(np.arange(self.n_samples), idx_full)] = labels.max() + 1
        # some labels will now not exist anymore as replaced by len(weights) - 1.
        # this means that the associated weights should be 0.
        # by doing np.bincount(labels)*weights this will be taken into account
        self.train_labels = labels
        self.idx_full = idx_full
    if self.store_location is not None:
        if (
            not os.path.exists(
                os.path.join(self.store_location, "train_labels.npy")
            )
            or self.force_recompute_indices
        ):
            os.makedirs(self.store_location, exist_ok=True)
            if self.nnz is not None:
                np.save(os.path.join(self.store_location, "nnz.npy"), self.nnz)
            np.save(
                os.path.join(self.store_location, "train_labels.npy"),
                self.train_labels,
            )
            np.save(
                os.path.join(self.store_location, "idx_full.npy"), self.idx_full
            )
            if self.test_idx is not None:
                np.save(
                    os.path.join(self.store_location, "test_idx.npy"), self.test_idx
                )
            if self.valid_idx is not None:
                np.save(
                    os.path.join(self.store_location, "valid_idx.npy"),
                    self.valid_idx,
                )
            listToFile(
                self.test_datasets,
                os.path.join(self.store_location, "test_datasets.txt"),
            )
        else:
            self.nnz = (
                np.load(os.path.join(self.store_location, "nnz.npy"), mmap_mode="r")
                if os.path.exists(os.path.join(self.store_location, "nnz.npy"))
                else None
            )
            self.train_labels = np.load(
                os.path.join(self.store_location, "train_labels.npy")
            )
            self.idx_full = np.load(
                os.path.join(self.store_location, "idx_full.npy"), mmap_mode="r"
            )
            self.test_idx = (
                np.load(os.path.join(self.store_location, "test_idx.npy"))
                if os.path.exists(os.path.join(self.store_location, "test_idx.npy"))
                else None
            )
            self.valid_idx = (
                np.load(os.path.join(self.store_location, "valid_idx.npy"))
                if os.path.exists(
                    os.path.join(self.store_location, "valid_idx.npy")
                )
                else None
            )
            self.test_datasets = fileToList(
                os.path.join(self.store_location, "test_datasets.txt")
            )
            print("loaded from store")
    print(f"done setup, took {time.time() - start_time:.2f} seconds")
    return self.test_datasets

def test_dataloader(self):
    """
    Create the test DataLoader with sequential sampling.

    Returns:
        DataLoader | List: Test DataLoader, or empty list if no test split.
    """
    return (
        DataLoader(
            self.dataset, sampler=SequentialSampler(self.test_idx), **self.kwargs
        )
        if self.test_idx is not None
        else []
    )

def train_dataloader(self, **kwargs):
    """
    Create the training DataLoader with weighted random sampling.

    Uses LabelWeightedSampler for class-balanced sampling when weight_scaler > 0
    and clss_to_weight is specified. Otherwise uses RankShardSampler for
    distributed training without weighting.

    Args:
        **kwargs: Additional arguments passed to DataLoader, overriding defaults.

    Returns:
        DataLoader: Training DataLoader instance.

    Raises:
        ValueError: If setup() has not been called.
    """
    if len(self.clss_to_weight) > 0 and self.weight_scaler > 0:
        try:
            print("Setting up the parallel train sampler...")
            # Create the optimized parallel sampler
            print(f"Using {self.sampler_workers} workers for class indexing")
            train_sampler = LabelWeightedSampler(
                labels=self.train_labels,
                weight_scaler=self.weight_scaler,
                num_samples=int(self.n_samples_per_epoch),
                element_weights=self.nnz,
                replacement=self.replacement,
                n_workers=self.sampler_workers,
                chunk_size=self.sampler_chunk_size,
                store_location=self.store_location,
                force_recompute_indices=self.force_recompute_indices,
                curiculum=self.curiculum,
            )
        except ValueError as e:
            raise ValueError(str(e) + " Have you run `datamodule.setup()`?")
        dataset = None
    else:
        dataset = Subset(self.dataset, self.idx_full)
        train_sampler = RankShardSampler(len(dataset), start_at=self.start_at)
    current_loader_kwargs = kwargs.copy()
    current_loader_kwargs.update(self.kwargs)
    return DataLoader(
        self.dataset if dataset is None else dataset,
        sampler=train_sampler,
        **current_loader_kwargs,
    )

def val_dataloader(self):
    """
    Create the validation DataLoader.

    Returns:
        DataLoader | List: Validation DataLoader, or empty list if no validation split.
    """
    return (
        DataLoader(
            Subset(self.dataset, self.valid_idx),
            **self.kwargs,
        )
        if self.valid_idx is not None
        else []
    )