Autoimputation

import warnings

warnings.filterwarnings("ignore")

import logging

logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
logging.getLogger("pytorch_tabular").setLevel(logging.ERROR)
logging.getLogger("joblib").setLevel(logging.ERROR)

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from sklearn.datasets import load_diabetes

pd.set_option("display.width", 600)
pd.set_option("display.max_columns", 10)
pd.set_option("display.expand_frame_repr", False)

from microimpute.comparisons.autoimpute import autoimpute
from microimpute.models import OLS, QuantReg, QRF, Matching
from microimpute.visualizations.comparison_plots import method_comparison_results

# Load the diabetes dataset
diabetes = load_diabetes()
diabetes_data = pd.DataFrame(diabetes.data, columns=diabetes.feature_names)

# Display the first few rows to understand the data structure
diabetes_data.head()

# Split the data into donor and receiver portions
donor_indices = np.random.choice(
    len(diabetes_data), size=int(0.7 * len(diabetes_data)), replace=False
)
receiver_indices = np.array(
    [i for i in range(len(diabetes_data)) if i not in donor_indices]
)

donor_data = diabetes_data.iloc[donor_indices].reset_index(drop=True)
receiver_data = diabetes_data.iloc[receiver_indices].reset_index(drop=True)

# Define which variables we'll use as predictors and which we want to impute
predictors = ["age", "sex", "bmi", "bp"]
imputed_variables = ["s1", "s4"]

# For demonstration purposes, we'll remove the variables we want to impute from the receiver dataset
receiver_data_without_targets = receiver_data.drop(columns=imputed_variables)

print(f"Donor data shape: {donor_data.shape}")
print(f"Receiver data shape: {receiver_data_without_targets.shape}")
print(f"Predictors: {predictors}")
print(f"Variables to impute: {imputed_variables}")

warnings.filterwarnings("ignore")

# Run the autoimpute process
results = autoimpute(
    donor_data=donor_data,
    receiver_data=receiver_data_without_targets,
    predictors=predictors,
    imputed_variables=imputed_variables,
    models=[OLS, QuantReg, QRF, Matching],  # MDN model excluded for efficiency
    tune_hyperparameters=False,
    k_folds=3,
)

print(
    f"Shape of receiver data before imputation: {receiver_data_without_targets.shape} \nShape of receiver data after imputation: {results.receiver_data.shape}"
)

# Examine the comparative performance of different imputation methods
print("Cross-validation results for different imputation methods:")
for model, metric_dict in results.cv_results.items():
    print(f"\nModel: {model}")
    print(
        f"quantile loss results: {metric_dict.get('quantile_loss').get('mean_test'):.4f}"
    )

# Identify which method was selected as the best performer
print(f"Best performing method: {results.fitted_models['best_method']}")

# Extract the quantiles used in the evaluation
comparison_viz = method_comparison_results(
    data=results.cv_results,
    metric="quantile_loss",
    data_format="wide",
)
fig = comparison_viz.plot(
    title="Autoimpute Method Comparison",
    show_mean=True,
)
fig.show()

# Examine imputed values (these were imputed for q=0.5 by default)
median_imputations = results.imputations[
    "best_method"
]  # Extract the best imputations with the "best_method" key
print("Median imputed values:")
median_imputations.head()

# Look at the full receiver dataset with imputed values integrated
print("Receiver dataset with imputed values:")
results.receiver_data.head()

# Visualize comparison between actual and imputed values
for var in imputed_variables:
    fig = go.Figure()

    # Plot actual values
    fig.add_trace(
        go.Scatter(
            x=receiver_data.index,
            y=receiver_data[var],
            mode="markers",
            name="Actual values",
            marker=dict(color="blue", size=8),
        )
    )

    # Plot imputed values
    fig.add_trace(
        go.Scatter(
            x=results.receiver_data.index,
            y=results.receiver_data[var],
            mode="markers",
            name="Imputed values",
            marker=dict(color="red", size=8),
        )
    )

    # Customize the plot appearance
    fig.update_layout(
        title=f"Comparison of actual vs imputed values for {var}",
        xaxis_title="Sample Index",
        yaxis_title=f"{var} Value",
        legend_title="Type",
        hovermode="closest",
    )

    fig.show()

from microimpute.models import *

# Specify a custom subset of models to evaluate
custom_models = [QRF, OLS, Matching]

# Specify custom hyperparameters for some models
custom_hyperparameters = {
    "QRF": {"n_estimators": 200, "max_depth": 10},
    "Matching": {"constrained": True},
}

# Then simply run autoimpute with custom models and hyperparameters
advanced_results = autoimpute(
    donor_data=donor_data,
    receiver_data=receiver_data_without_targets,
    predictors=predictors,
    imputed_variables=imputed_variables,
    models=custom_models,
    hyperparameters=custom_hyperparameters,
    k_folds=3,
)

advanced_results.imputations["best_method"]

warnings.filterwarnings("ignore")

# Run the autoimpute process
results = autoimpute(
    donor_data=donor_data,
    receiver_data=receiver_data_without_targets,
    predictors=predictors,
    imputed_variables=imputed_variables,
    models=[OLS, QuantReg, QRF, Matching],  # MDN model excluded for efficiency
    tune_hyperparameters=False,
    impute_all=True,
    k_folds=3,
)

print(f"Imputation results available for models: {results.imputations.keys()}")
print(
    f"The best performing model is: {results.fitted_models['best_method'].__class__.__name__}"
)