data_processing.py

import json
import os
from collections import Counter
from functools import reduce
import numpy as np
import pandas as pd
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer, SimpleImputer
from sklearn.linear_model import LinearRegression

from src.cleanup import (
    fix_faq_values,
    get_exam_date_column,
    get_patient_id_column,
    get_visit_id_column,
    handle_truncated_measurements,
    normalize_sex,
    setup_preprocessing,
)
from src.feature_engineering import Feature
from src.helpers import check_missing_data, merge_dataframes_list
from src.model_training import ensure_distribution, train_with_folds


def get_target(features_json):
    with open(features_json, "r") as features_file:
        features_config = json.load(features_file)["features"]
    features_config_target = features_config["target"]
    return Feature(
        features_config_target["name"],
        features_config_target["dtype"],
        is_composite=False,
        acceptable_values=features_config_target.get("acceptable_values"),
        is_target=True,
    )


def add_features(config: dict):
    features = list[Feature]()
    for field, value in config["simple"].items():
        features.append(Feature(field, value["dtype"], is_composite=False))

    for field, value in config["composite"].items():
        field_feat_refs = list[Feature]()
        for field_name in value["fields"]:
            if [f.name for f in features].count(field_name) > 0:
                item = [f for f in features if f.name == field_name][0]
                field_feat_refs.append(item)
            else:
                raise ValueError(
                    f"Field {field_name} not found in prior features list."
                )
        features.append(
            Feature(
                field,
                value["dtype"],
                is_composite=True,
                fields=field_feat_refs,
                operation=value["operation"],
                acceptable_values=value.get("acceptable_values"),
            )
        )

    value = config["target"]
    features.append(
        Feature(
            value["name"],
            d_type=value["dtype"],
            is_composite=False,
            acceptable_values=value.get("acceptable_values"),
            is_target=True,
        )
    )

    return features


OUTPUT_FILE_PATH = "data/processed/dataset_cleaned.csv"
CONFIG_PATH = "data/maps"
FEATURES_JSON_PATH = os.path.join(CONFIG_PATH, "features.json")
target = get_target(FEATURES_JSON_PATH)


def preprocess(data_json: str, features_json: str):
    with open(data_json, "r") as data_file:
        data_config = json.load(data_file)["datasets"]
    with open(features_json, "r") as features_file:
        features_config = json.load(features_file)["features"]

    features = add_features(features_config)
    dataframes = list[pd.DataFrame]()

    for dataset in data_config:
        dataset_name = str(dataset["name"])
        dataset_file_paths = dataset["paths"]
        dataset_fields = {**dataset["fields"], **dataset["target"]}
        feature_field_names = [
            dataset_fields[feature.name]
            for feature in features
            if not feature.is_composite
        ]
        dataframe = pd.DataFrame()
        sub_dataframes = list[pd.DataFrame]()
        ptid = get_patient_id_column(dataset_name, data_config)
        vid = get_visit_id_column(dataset_name, data_config)
        exam_date = get_exam_date_column(dataset_name, data_config)

        for path in dataset_file_paths:
            feature_names = feature_field_names.copy()

            for sub_frame in sub_dataframes:
                for field in feature_field_names:
                    if (
                        field in sub_frame.columns
                        and field != ptid
                        and field != vid
                        and field != exam_date
                    ):
                        feature_names.remove(field)
            df = pd.read_csv(
                f"data/raw/{path}",
                usecols=lambda x: x in feature_names,
                low_memory=False,
            )

            sub_dataframes.append(df)

        # Identify common features across all datasets
        common_features = list(
            reduce(
                lambda x, y: set(x) & set(y),
                [df.columns for df in sub_dataframes],
            )
        )

        dataframe = merge_dataframes_list(sub_dataframes, common_features)
        dataframe = dataframe.rename(
            mapper={v: k for k, v in dataset_fields.items()}, axis=1, errors="raise"
        )

        dataframe = handle_truncated_measurements(dataframe)
        dataframe["Dataset"] = dataset_name

        # Normalize values for Sex to 0 and 1
        dataframe = normalize_sex(dataframe)
        # fix faq values
        dataframe = fix_faq_values(dataframe)
        # replace -4.0 with NaN
        dataframe = dataframe.replace(to_replace=-4.0, value=np.nan)
        # replace HasDepression-specific 9 with NaN
        dataframe["HasDepression"] = dataframe["HasDepression"].replace(
            to_replace=[9], value=np.nan
        )
        dataframes.append(dataframe)

    train_set = [df for df in dataframes if (df["Dataset"] == "NACC").any()][0]
    # combined_dataset = pd.concat(dataframes, ignore_index=True)
    combined_dataset = train_set
    dataset_less_target = combined_dataset.drop(target.name, axis=1)
    dataset_less_target = dataset_less_target.drop(
        ["PatientId", "VisitId", "VisitDate", "Dataset"], axis=1
    )

    # Impute missing boolean values
    bool_features = ["HasDepression", "HasStroke", "HasCircumstance"]
    imputer = IterativeImputer(
        estimator=LinearRegression(), random_state=42, initial_strategy="most_frequent"
    )
    print("Imputing missing values for boolean features.")
    imputed_dataset = imputer.fit_transform(dataset_less_target)
    for col in bool_features:
        dataset_less_target[col] = imputed_dataset[
            :, dataset_less_target.columns.get_loc(col)
        ]
    dataset_less_target[bool_features] = dataset_less_target[bool_features].astype(bool)

    # Impute missing values
    numerical_features = dataset_less_target.select_dtypes(include=["number"]).columns
    imputer = SimpleImputer(strategy="mean")
    print("Imputing missing values for numerical features.")
    dataset_less_target[numerical_features] = np.round(
        imputer.fit_transform(dataset_less_target[numerical_features])
    )

    preproc_transformer = setup_preprocessing(features=features)
    preproc_transformer.set_output(transform="pandas")
    preproc_dataset = preproc_transformer.fit_transform(dataset_less_target)

    for col in preproc_dataset.columns:
        dataset_less_target[col] = preproc_dataset[col]

    combined_dataset[dataset_less_target.columns] = dataset_less_target
    columns_post_transform = preproc_transformer.get_feature_names_out()
    columns_post_transform = np.append(columns_post_transform, target.name)

    transformed_dataset = pd.DataFrame(combined_dataset, columns=columns_post_transform)
    transformed_dataset = ensure_distribution(transformed_dataset, target)

    transformed_dataset.to_csv(OUTPUT_FILE_PATH, index=False)

    print("All steps for data cleaning, class balancing, and imputation are complete.")
    print(
        f"A file containing the compilation results was written to {OUTPUT_FILE_PATH}"
    )

    check_missing_data(OUTPUT_FILE_PATH)
    return transformed_dataset


if __name__ == "__main__":
    data_json_path = os.path.join(CONFIG_PATH, "datasets.json")
    trainable = preprocess(data_json_path, FEATURES_JSON_PATH)
    print("Preprocessing complete.\nStarting model training...")
    train_with_folds(trainable, target)