Quantile Regression Forest (QRF) imputation

# Import necessary libraries
import warnings

warnings.filterwarnings("ignore")

import logging

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

import pandas as pd
import numpy as np
import plotly.express as px
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)

# Import Microimpute tools
from microimpute.utils.data import preprocess_data
from microimpute.evaluations import cross_validate_model
from microimpute.models import QRF
from microimpute.config import QUANTILES
from microimpute.visualizations import model_performance_results

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

# Display the first few rows of the dataset
df.head()

# Define variables for the model
predictors = ["age", "sex", "bmi", "bp"]
imputed_variables = [
    "s1",
    "s4",
]  # We'll impute 's1' (total serum cholesterol) and 's4' (total cholesterol/HDL ratio)

# Create a subset with only needed columns
diabetes_df = df[predictors + imputed_variables]

# Display summary statistics
diabetes_df.describe()

# Split data into training and testing sets
X_train, X_test = preprocess_data(diabetes_df)

# Let's see how many records we have in each set
print(f"Training set size: {X_train.shape[0]} records")
print(f"Testing set size: {X_test.shape[0]} records")

# Create a copy of the test set with missing values
X_test_missing = X_test.copy()

# Store the actual values for later comparison
actual_values = X_test_missing[imputed_variables].copy()

# Remove the values to be imputed
X_test_missing[imputed_variables] = np.nan

X_test_missing.head()

# Define quantiles we want to model
# We'll use the default quantiles from the config module
print(f"Modeling these quantiles: {QUANTILES}")

# Initialize the QRF imputer with some custom parameters
# You can customize the random forest by passing additional parameters
qrf_imputer = QRF()

# Fit the model with our training data
# This trains a quantile regression forest model
fitted_qrf_imputer = qrf_imputer.fit(
    X_train,
    predictors,
    imputed_variables,
    n_estimators=100,
    min_samples_leaf=5,
)

# Impute values in the test set
# This uses the trained QRF model to predict missing values at specified quantiles
imputed_values = fitted_qrf_imputer.predict(X_test_missing, QUANTILES)

# Display the first few imputed values at the median (0.5 quantile)
imputed_values[0.5].head()

# Define your quantiles
quantiles = list(imputed_values.keys())

# Convert imputed_values dict to a 3D array: (n_samples, n_quantiles)
pred_matrix = np.stack([imputed_values[q].values.flatten() for q in quantiles], axis=1)

# Actual values flattened
actual = actual_values.values.flatten()

# Compute absolute error matrix: shape (n_samples, n_quantiles)
abs_error = np.abs(pred_matrix - actual[:, None])

# Find index of closest prediction for each sample
closest_indices = abs_error.argmin(axis=1)

# Select the closest predictions
closest_predictions = np.array(
    [pred_matrix[i, idx] for i, idx in enumerate(closest_indices)]
)

# Wrap as DataFrame for plotting
closest_df = pd.DataFrame(
    {
        "Actual": actual,
        "ClosestPrediction": closest_predictions,
    }
)

# Extract median predictions for evaluation
median_predictions = imputed_values[0.5]

# Create a scatter plot comparing actual vs. imputed values
min_val = min(actual_values.min().min(), median_predictions.min().min())
max_val = max(actual_values.max().max(), median_predictions.max().max())

# Create the scatter plot
fig = px.scatter(
    closest_df,
    x="Actual",
    y="ClosestPrediction",
    opacity=0.7,
    title="Comparison of actual vs. imputed values using QRF",
)

# Add the diagonal line (perfect prediction line)
fig.add_trace(
    go.Scatter(
        x=[min_val, max_val],
        y=[min_val, max_val],
        mode="lines",
        line=dict(color="red", dash="dash"),
        name="Perfect Prediction",
    )
)

# Update layout
fig.update_layout(
    xaxis_title="Actual values",
    yaxis_title="Imputed values",
    width=750,
    height=600,
    template="plotly_white",
    margin=dict(l=50, r=50, t=80, b=50),  # Adjust margins
)

fig.show()

# Compare predictions at different quantiles for the first 5 records
quantiles_to_show = QUANTILES
comparison_df = pd.DataFrame(index=range(5))

# Add actual values
comparison_df["Actual"] = actual_values.iloc[:5, 0].values

# Add quantile predictions
for q in quantiles_to_show:
    comparison_df[f"Q{int(q * 100)}"] = imputed_values[q].iloc[:5, 0].values

comparison_df

# Create a prediction interval plot for the first 10 records
# Number of records to plot
n_records = 10

# Prepare data for plotting
records = list(range(n_records))
actuals = actual_values.iloc[:n_records, 0].values
medians = imputed_values[0.5].iloc[:n_records, 0].values
q30 = imputed_values[0.3].iloc[:n_records, 0].values
q70 = imputed_values[0.7].iloc[:n_records, 0].values
q10 = imputed_values[0.1].iloc[:n_records, 0].values
q90 = imputed_values[0.9].iloc[:n_records, 0].values

# Create the base figure
fig = go.Figure()

# Add 80% prediction interval (Q10-Q90)
for i in range(n_records):
    fig.add_trace(
        go.Scatter(
            x=[i, i],
            y=[q10[i], q90[i]],
            mode="lines",
            line=dict(width=10, color="rgba(173, 216, 230, 0.3)"),
            hoverinfo="none",
            showlegend=False,
        )
    )

# Add 40% prediction interval (Q30-Q70)
for i in range(n_records):
    fig.add_trace(
        go.Scatter(
            x=[i, i],
            y=[q30[i], q70[i]],
            mode="lines",
            line=dict(width=10, color="rgba(70, 130, 180, 0.5)"),
            hoverinfo="none",
            showlegend=False,
        )
    )

# Add actual values
fig.add_trace(
    go.Scatter(
        x=records,
        y=actuals,
        mode="markers",
        marker=dict(color="black", size=8),
        name="Actual",
    )
)

# Add median predictions
fig.add_trace(
    go.Scatter(
        x=records,
        y=medians,
        mode="markers",
        marker=dict(color="red", size=8),
        name="Median (Q50)",
    )
)

# Add dashed line for Q10
fig.add_trace(
    go.Scatter(
        x=[-1, -1],  # Dummy points for legend
        y=[0, 0],  # Dummy points for legend
        mode="lines",
        line=dict(color="rgba(173, 216, 230, 0.3)", width=10),
        name="80% PI (Q10-Q90)",
    )
)

# Add dashed line for Q30
fig.add_trace(
    go.Scatter(
        x=[-1, -1],  # Dummy points for legend
        y=[0, 0],  # Dummy points for legend
        mode="lines",
        line=dict(color="rgba(70, 130, 180, 0.5)", width=10),
        name="40% PI (Q30-Q70)",
    )
)

# Update layout with smaller width to fit in the book layout
fig.update_layout(
    title="QRF imputation prediction intervals",
    xaxis=dict(
        title="Data record index",
        showgrid=True,
        gridwidth=1,
        gridcolor="rgba(211, 211, 211, 0.7)",
    ),
    yaxis=dict(
        title="Total Serum Cholesterol (s1)",
        showgrid=True,
        gridwidth=1,
        gridcolor="rgba(211, 211, 211, 0.7)",
    ),
    width=750,
    height=600,
    template="plotly_white",
    margin=dict(l=50, r=50, t=80, b=50),  # Adjust margins
    legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99),
)

fig.show()

predictors = ["age", "sex", "bmi", "bp"]
imputed_variables = ["s1", "s4"]

# Run cross-validation on the same data set
qrf_results = cross_validate_model(QRF, diabetes_df, predictors, imputed_variables)

# Check if we have quantile loss results (for numerical variables)
if "quantile_loss" in qrf_results:
    print("Quantile loss results:")
    print(qrf_results["quantile_loss"]["results"])

# Plot the results for numerical variables
if "quantile_loss" in qrf_results:
    perf_results_viz = model_performance_results(
        results=qrf_results,
        model_name="QRF",
        method_name="Cross-validation quantile loss average",
    )
    fig = perf_results_viz.plot(
        title="QRF cross-validation performance",
    )
    fig.show()

# Check the hyperparameters with which QRF was initialized
print(fitted_qrf_imputer.models[imputed_variables[0]].qrf.get_params())

# To set specific hyperparameters pass them when fitting the model
fitted_qrf_imputer = qrf_imputer.fit(
    X_train=df,
    predictors=predictors,
    imputed_variables=imputed_variables,
    n_estimators=200,
    min_samples_leaf=10,
    max_depth=5,
)

print(fitted_qrf_imputer.models[imputed_variables[0]].qrf.get_params())

# Create a dataset with categorical variables
np.random.seed(42)

# Create synthetic categorical variables based on diabetes features
df_categorical = pd.DataFrame()
df_categorical["age"] = df["age"]
df_categorical["sex"] = df["sex"]
df_categorical["bmi"] = df["bmi"]
df_categorical["bp"] = df["bp"]
df_categorical["risk_level"] = pd.qcut(
    df["s1"],
    q=3,
    labels=["low", "medium", "high"],
).astype(str)

print("Categorical variable distribution:")
print(pd.Series(df_categorical["risk_level"]).value_counts())
print(f"\nPercentage distribution:")
print(pd.Series(df_categorical["risk_level"]).value_counts(normalize=True))
print(f"\nData types: {df_categorical.dtypes.to_dict()}")

# Split the categorical data for training and testing
X_train_cat, X_test_cat = preprocess_data(df_categorical)

print(f"\nTraining set size: {X_train_cat.shape[0]} records")
print(f"Testing set size: {X_test_cat.shape[0]} records")

# Fit QRF model for categorical imputation
predictors_cat = ["age", "sex", "bmi", "bp"]
imputed_variables_cat = ["risk_level"]

# Initialize and fit the QRF imputer
qrf_cat_imputer = QRF()
fitted_qrf_cat = qrf_cat_imputer.fit(X_train_cat, predictors_cat, imputed_variables_cat)

print("QRF model fitted for categorical variable imputation")

# Create test set with missing categorical values
X_test_cat_missing = X_test_cat.copy()
actual_cat_values = X_test_cat_missing[imputed_variables_cat].copy()
X_test_cat_missing[imputed_variables_cat] = np.nan

# Impute the categorical values
# For categorical variables, all quantiles return the same prediction
imputed_cat_values = fitted_qrf_cat.predict(X_test_cat_missing, [0.5])

# Evaluate categorical imputation accuracy
from sklearn.metrics import accuracy_score, confusion_matrix

# Get predictions and actual values
predicted = imputed_cat_values[0.5]["risk_level"].values
actual = actual_cat_values["risk_level"].values

# Calculate accuracy
accuracy = accuracy_score(actual, predicted)
print(f"Categorical imputation accuracy: {accuracy:.2%}")

# Create confusion matrix
conf_matrix = pd.DataFrame(
    confusion_matrix(actual, predicted),
    index=["Actual: low", "Actual: medium", "Actual: high"],
    columns=["Predicted: low", "Predicted: medium", "Predicted: high"],
)
print("\nConfusion matrix:")
print(conf_matrix)

# Run cross-validation for categorical variables
predictors_cat = ["age", "sex", "bmi", "bp"]
imputed_variables_cat = ["risk_level"]

qrf_categorical_results = cross_validate_model(
    QRF, df_categorical, predictors_cat, imputed_variables_cat
)

# Display results
print("Categorical imputation cross-validation results (log loss):")
print(f"Mean train log loss: {qrf_categorical_results['log_loss']['mean_train']:.4f}")
print(f"Mean test log loss: {qrf_categorical_results['log_loss']['mean_test']:.4f}")

# Plot the categorical imputation performance
cat_perf_results_viz = model_performance_results(
    results=qrf_categorical_results,
    model_name="QRF",
    method_name="Cross-validation log loss average",
    metric="log_loss",
)
fig = cat_perf_results_viz.plot(
    title="QRF categorical imputation cross-validation performance",
)
fig.show()