data_handlers ¶

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def encode_dataframe(self, dataframe: pl.DataFrame) -> dict[str, torch.Tensor]:
    """Encode the dataframe columns using the configured encoders.

    Takes a polars DataFrame and encodes each column using its corresponding encoder
    from self.encoders.

    Args:
        dataframe: Polars DataFrame containing the columns to encode

    Returns:
        Dict mapping column names to their encoded tensors. The exact shape of each
        tensor depends on the encoder used for that column.

    Example:
        >>> df = pl.DataFrame({"dna_seq": ["ACGT", "TGCA"], "labels": [1, 2]})
        >>> encoded = dataset_loader.encode_dataframe(df)
        >>> print(encoded["dna_seq"].shape)
        torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
    """
    return {col: self.encoders[col].encode_all(dataframe[col].to_list()) for col in dataframe.columns}

def get_all_items(self) -> tuple[dict, dict, dict]:
    """Get the full dataset as three separate dictionaries for inputs, labels and metadata.

    Returns:
        tuple[dict, dict, dict]: Three dictionaries containing:
            - Input dictionary mapping input column names to encoded input data
            - Label dictionary mapping label column names to encoded label data
            - Meta dictionary mapping meta column names to meta data

    Example:
        >>> handler = DatasetHandler(...)
        >>> input_dict, label_dict, meta_dict = handler.get_dataset()
        >>> print(input_dict.keys())
        dict_keys(['age', 'fare'])
        >>> print(label_dict.keys())
        dict_keys(['survived'])
        >>> print(meta_dict.keys())
        dict_keys(['passenger_id'])
    """
    input_data = self.encode_dataframe(self.data[self.input_columns])
    label_data = self.encode_dataframe(self.data[self.label_columns])
    meta_data = {key: self.data[key].to_list() for key in self.meta_columns}
    return input_data, label_data, meta_data

def get_all_items_and_length(self) -> tuple[tuple[dict, dict, dict], int]:
    """Get the full dataset as three separate dictionaries for inputs, labels and metadata, and the length of the data."""
    return self.get_all_items(), len(self.data)

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def load_csv_per_split(self, csv_path: str, split: int) -> pl.DataFrame:
    """Load the part of csv file that has the specified split value.

    Split is a number that for 0 is train, 1 is validation, 2 is test.
    """
    if "split" not in self.columns:
        raise ValueError("The category split is not present in the csv file")
    if split not in [0, 1, 2]:
        raise ValueError(
            f"The split value should be 0, 1 or 2. The specified split value is {split}",
        )
    return pl.scan_csv(csv_path).filter(pl.col("split") == split).collect()

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def add_split(self, *, force: bool = False) -> None:
    """Add a column specifying the train, validation, test splits of the data.

    An error exception is raised if the split column is already present in the csv file. This behaviour can be overriden by setting force=True.

    Args:
        force (bool): If True, the split column present in the csv file will be overwritten.
    """
    if ("split" in self.columns) and (not force):
        raise ValueError(
            "The category split is already present in the csv file. If you want to still use this function, set force=True",
        )
    # get relevant split columns from the dataset_manager
    split_input_data = self.select_columns(self.split_columns)

    # get the split indices
    train, validation, test = self.splitter.get_split_indexes(split_input_data)

    # add the split column to the data
    split_column = np.full(len(self.data), -1).astype(int)
    split_column[train] = 0
    split_column[validation] = 1
    split_column[test] = 2
    self.data = self.data.with_columns(pl.Series("split", split_column))

    if "split" not in self.columns:
        self.columns.append("split")

def apply_transformations(self) -> None:
    """Apply the transformation group to the data.

    Applies all transformations defined in self.transforms to their corresponding columns.
    Each column can have multiple transformations that are applied sequentially.
    """
    for column_name, transforms_list in self.transforms.items():
        for transform in transforms_list:
            transformed_data = transform.transform_all(self.data[column_name].to_list())

            if transform.add_row:
                new_rows = self.data.with_columns(
                    pl.Series(column_name, transformed_data),
                )
                self.data = pl.concat([self.data, new_rows], how="vertical")
            else:
                self.data = self.data.with_columns(
                    pl.Series(column_name, transformed_data),
                )

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def shuffle_labels(self, label_columns: list[str], seed: Optional[float] = None) -> None:
    """Shuffles the labels in the data."""
    # set the np seed
    np.random.seed(seed)

    for key in label_columns:
        self.data = self.data.with_columns(
            pl.Series(key, np.random.permutation(list(self.data[key]))),
        )