Datamanager

def complete_dataset_initialization(
    self, controller_kwargs: Dict[str, Any]
) -> None:
    """Finalizes initialization of the DataManager's dataset controller

    Args:
        controller_kwargs: arguments for the controller

    Raises:
        FedbiomedError: if `_data_manager_instance` is not initialized
        FedbiomedError: if there is a problem completing dataset initialization
    """
    if not self._data_manager_instance:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Data manager instance is not initialized. "
            f"Please call `load()` first."
        )

    try:
        self._dataset.complete_initialization(
            controller_kwargs,
            _dm_to_format[self._data_manager_instance.__class__],
        )
    except FedbiomedError as e:
        raise e
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Unable to complete dataset initialization."
        ) from e

def extend_loader_args(self, extension: Optional[Dict]) -> None:
    """Extends the class' loader arguments

    Extends the class's `_loader_arguments` attribute with additional key-values from
    the `extension` argument. If a key already exists in the `_loader_arguments`, then
    it is not replaced.

    Args:
        extension: the mapping used to extend the loader arguments
    """
    if extension:
        self._loader_arguments.update(
            {
                key: value
                for key, value in extension.items()
                if key not in self._loader_arguments
            }
        )

def load(self, tp_type: TrainingPlans) -> None:
    """Loads proper DataManager based on given TrainingPlan and
    `dataset`, `target` attributes.

    Args:
        tp_type: Enumeration instance of TrainingPlans that stands for type of training plan.

    Raises:
        FedbiomedError: unknown training plan type

    """
    if tp_type in _tp_to_datamanager:
        self._data_manager_instance = _tp_to_datamanager[tp_type](
            dataset=self._dataset, **self._loader_arguments
        )
    else:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Unknown training plan type, "
            "cannot instantiate data manager."
        )

def load_state(self, state: Dict):
    """Loads state of the data loader


    It currently keep only testing index, training index and test ratio
    as state.

    Args:
        state: Object containing data loader state.
    """
    self._testing_index = state.get("testing_index", [])
    self._training_index = state.get("training_index", [])
    self._test_ratio = state.get("test_ratio", None)

def save_state(self) -> Dict:
    """Gets state of the data loader.

    Returns:
        A Dict containing data loader state.
    """

    data_manager_state = {}
    data_manager_state["training_index"] = self._training_index
    data_manager_state["testing_index"] = self._testing_index
    data_manager_state["test_ratio"] = self._test_ratio
    return data_manager_state

def split(
    self,
    test_ratio: float,
    test_batch_size: Optional[int],
    is_shuffled_testing_dataset: bool = False,
) -> Tuple[Optional[DataLoader], Optional[DataLoader]]:
    """Split Dataset into train and validation dataloaders.

    Args:
        test_ratio: Split ratio for validation set ratio. Rest of the samples will be used for training
        test_batch_size: Batch size to use for testing subset
        is_shuffled_testing_dataset: if True, randomly select different samples for the testing
            subset at each execution. If False, reuse previous split when possible.
    Raises:
        FedbiomedError: Arguments bad format
        FedbiomedError: Cannot get number of samples from dataset

    Returns:
        train_loader: DataLoader for training subset. `None` if the `test_ratio` is `1`
        test_loader: DataLoader for validation subset. `None` if the `test_ratio` is `0`
    """
    # No need to check is_shuffled_testing_dataset, any argument can be interpreted as bool

    # Check the type of argument test_batch_size
    if not isinstance(test_batch_size, int) and test_batch_size is not None:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: The argument `test_batch_size` should be "
            f"type `int` or `None` not {type(test_batch_size)}"
        )

    # Check the argument `ratio` is of type `float`
    if not isinstance(test_ratio, (float, int)):
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: The argument `test_ratio` should be "
            f"type `float` or `int` not {type(test_ratio)}"
        )

    # Check ratio is valid for splitting
    if test_ratio < 0 or test_ratio > 1:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: The argument `test_ratio` should be "
            f"equal or between 0 and 1, not {test_ratio}"
        )

    # Wrap dataset in framework specific class if needed
    framework_dataset = self._dataset_wrapper(self._dataset)

    try:
        samples = len(framework_dataset)
    except AttributeError as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Can not get number of samples from "
            f"{str(self._dataset)} due to undefined attribute, {str(e)}"
        ) from e
    except TypeError as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Can not get number of samples from "
            f"{str(self._dataset)}, {str(e)}"
        ) from e

    if self._test_ratio != test_ratio and self._test_ratio is not None:
        if not is_shuffled_testing_dataset:
            logger.info(
                "`test_ratio` value has changed: this will change the testing dataset"
            )
        is_shuffled_testing_dataset = True

    _is_loading_failed: bool = False
    # Calculate number of samples for train and validation subsets
    test_samples = math.floor(samples * test_ratio)
    train_samples = samples - test_samples

    if self._testing_index and not is_shuffled_testing_dataset:
        try:
            self._load_indexes(
                framework_dataset, self._training_index, self._testing_index
            )
            _is_loading_failed = False
        except IndexError:
            _is_loading_failed = True

    need_new_split = (
        not self._testing_index or is_shuffled_testing_dataset or _is_loading_failed
    )

    if need_new_split:
        if self._loader_arguments.get("shuffle", True):
            # Random split (shuffled)
            self._subset_train, self._subset_test = self._random_split(
                framework_dataset,
                [train_samples, test_samples],
            )
        else:
            # Deterministic split (no shuffle) — preserve original order
            all_indices = list(range(samples))
            train_indices = all_indices[:train_samples]
            test_indices = all_indices[train_samples : train_samples + test_samples]

            self._subset_train = self._subset_class(
                framework_dataset, train_indices
            )
            self._subset_test = (
                self._subset_class(framework_dataset, test_indices)
                if test_samples > 0
                else None
            )

        self._training_index = list(self._subset_train.indices)
        self._testing_index = (
            list(self._subset_test.indices) if self._subset_test is not None else []
        )

    if not test_batch_size and self._subset_test is not None:
        test_batch_size = len(self._subset_test)

    self._test_ratio = test_ratio

    loaders = (
        self._subset_loader(self._subset_train, **self._loader_arguments),
        self._subset_loader(self._subset_test, batch_size=test_batch_size),
    )

    return loaders

def subset_test(self) -> Optional[_SkLearnSubset]:
    """Gets validation subset of the dataset.

    Returns:
        Validation subset
    """

    return self._subset_test

def subset_train(self) -> Optional[_SkLearnSubset]:
    """Gets train subset of the dataset.

    Returns:
        Train subset
    """
    return self._subset_train

def load_all_samples(self) -> PytorchDataLoader:
    """Loading all samples as PyTorch DataLoader without splitting.

    Returns:
        Dataloader for entire datasets. `DataLoader` arguments will be retrieved from the `**kwargs` which
            is defined while initializing the class
    """
    torch_dataset = _DatasetWrapper(self._dataset)
    return self._create_data_loader(torch_dataset, **self._loader_arguments)  # type: ignore

def subset_test(self) -> Optional[_TorchSubset]:
    """Gets validation subset of the dataset.

    Returns:
        Validation subset
    """

    return self._subset_test

def subset_train(self) -> Optional[_TorchSubset]:
    """Gets train subset of the dataset.

    Returns:
        Train subset
    """
    return self._subset_train