MICE Data Imputation Model | Code Documentation¶

Description¶

Performs the data imputation model based on the MICE (Multiple Imputation by Chained Equations) algorithm method using the IterativeImputer auxiliary function from the sklearn.impute package.

The function receives a dataset with missing numeric and categorical registers, treats them using the MICE algorithm method and returns a version of the original dataset with imputed values ​​and a comparative statistical analysis table of both original and imputed dataset.

IterativeImputer MICE method¶

The Multiple Imputation by Chained Equations (MICE) method is an iterative approach to filling in missing values using predictive models. It models each feature with missing values as a function of other features, and uses that estimate for imputation. It does so in an iterated round-robin fashion: at each step, a feature column is designated as output y and the other feature columns are treated as inputs X. A regressor is fit on (X, y) for known y. Then, the regressor is used to predict the missing values of y. This is done for each feature in an iterative fashion, and then is repeated for max_iter imputation rounds. The results of the final imputation round are returned.

The IterativeImputer MICE method from scikit-learn is ruled by the following steps:

1. Initialize Missing Values

• Missing values are initially replaced with a simple estimate, such as the column mean, median or most frequent value.

2. Define an Imputation Strategy

• For each column with missing values, a regression model is defined.
• This model is trained to predict the missing values using the other columns as predictors.

3. Iterate Over Columns

• For each column with missing values, the dataset is split into training data (non-missing values) and test data (missing values).
• A regression model is trained on the available values and then used to predict the missing values.
• The imputed values are updated and used in the next iteration for further refinement.

4. Repeat Until Convergence

• This process is repeated for n_iter iterations (default = 10).
• At each iteration, the imputed values are refined as the models adapt better to the data.
• The final imputed values are those obtained after the last iteration.

Arguments¶

df: object, mandatory

The original dataset which can contain missing values or not and both numerial and categorical variables.

interest_vars: list, default = None

List formed by the name of the columns of the only numerical and categorical variables that will be used by the algorith to make the predictions and worn in the statistical analysis table. The list can contain the name of columns with no missing values - it will be used just as another parameter to incorporate the predictions. If "None", every column of the dataset will be used and analysed.

show_results: bool, default = True

Boolean argument for the user to choose if the function will make and show the statistical analysis of the interest variables on both original and imputed datasets.

random_state (sklearn): int, RandomState instance or None, default = None

The seed of the pseudo random number generator to use. Randomizes selection of estimator features if n_nearest_features is not None, the imputation_order if random, and the sampling from posterior if sample_posterior=True. Use an integer for determinism.

n: int, default = 5 Number of imputations (imputed datasets).

max_iter (sklearn): int, default = 10

Maximum number of imputation rounds in a single dataset to perform before returning the imputations computed during the final round. A round is a single imputation of each feature with missing values.

initial_strategy (sklearn): {'mean', 'median', 'most_frequent', 'constant'}, default = 'most_frequent'

Which strategy to use to initialize the missing values.

prefix_sep (pandas): str, default = '!'

If appending prefix, separator/delimiter to use in the binary transformation of categorical variables.

force_binary: bool, default = False

If True, forces binary outcomes on: (i) multi-class categorical variables (applied on their one-hot dummy blocks before reconstruction), and (ii) single-column binary categorical variables (e.g., already coded as 0/1). The threshold used is defined by cutoff (global) or cutoff_map (per-variable).

cutoff: float, default = 0.5

Global decision threshold for binarizing categorical variables when force_binary=True. Values >= cutoff become 1, values < cutoff become 0. Can be overridden per variable via cutoff_map.

cutoff_map: dict[str, float] or None, default = None

Optional per-variable threshold overrides when force_binary=True. Use a dictionary like {"IS_FEMALE": 0.4, "IS_HOSPITAL_DEATH": 0.9}. Variables not present in this mapping fall back to the global cutoff (default 0.5).

zero_row_strategy: {'highest_probability', 'nan', 'first', 'force_class'}, default = 'highest_probability'

Strategy used only for multi-class categorical variables when, after binarization (or due to all zeros), a row has all dummy values equal to 0:

• 'highest_probability': selects the class with the highest raw (pre-binarization) value and assigns it (even if below cutoff).
• 'nan': assigns NaN to that categorical variable.
• 'first': picks the first dummy column of the group (deterministic tie-break).
• 'force_class': forces a specific class provided via zero_row_force.

zero_row_force: str or dict[str, str], default = None

Class to be assigned when zero_row_strategy='force_class'.

• If a string is provided, the same class is used for all affected variables.
• If a dict is provided, map variable → class (e.g., {'CITY': 'NY', 'PLAN': 'gold'}). The forced class must exist among the original classes of each variable; otherwise, the function raises a clear error.

tie_strategy: {'highest_probability', 'first', 'nan', 'force_class'}, default = 'highest_probability'

Strategy to resolve ties above the cutoff (i.e., when force_binary=True and more than one class in a one-hot block is ≥ cutoff in the same row):

• 'highest_probability': break ties using the raw (pre-binarization) probabilities and pick the argmax;
• 'first': pick the first dummy column in the group among those equal to 1;
• 'nan': assign NaN to that categorical variable;
• 'force_class': assign a specific class provided via tie_force.

tie_force: str or dict[str, str], default = None

Class to be assigned when tie_strategy='force_class'.

If a string is provided, the same class is used for all tied variables.

If a dict is provided, map variable → class (e.g., {'CITY': 'RJ', 'PLAN': 'gold'}). The forced class must exist among the original classes of each variable; otherwise, the function raises a clear error.

imputer_kwargs_extra (sklearn): dict or None Additional kwargs to pass to IterativeImputer.

DISCLAIMER: Some arguments used in this MICE function are imported from other python packages as indicated next to the name of the argument.

Outputs¶

imputed_datasets : object

Zip file with "n" imputed datasets.

pooled_df: object

Pooled version of the "n" imputed datasets.

stats_df: object

Stat analysis dataset.

Statistical analysis table¶

VARIABLE:

Name of the column. If the column is categorical, it will show the name of the class followed by the name of the column.

MISSINGS:

Shows the percentage of missing values in the original dataset column.

BEFORE IMPUTATION:

For numerical variables, it will show the IQR (Interquartile Range) of the columns registers distributions before the imputation (original data) as "median value (Q1,Q3)". For categorical variables, it will show the total number of registers of the class and the percentage that it represents of the column before the imputation (original data) as "total count (percentage)"

AFTER IMPUTATION:

For numerical variables, it will show the IQR (Interquartile Range) of the columns registers distributions after the imputation (imputed data) as "median value (Q1,Q3)". For categorical variables, it will show the total number of registers of the class and the percentage that it represents of the column after the imputation (imputed dataset) as "total count (percentage)"

P-VALUE:

Statistical analysis that compares original and imputed dataset value distributions.

IMPUTATION VARIANCE:

For each imputed column, mesure the variance of imputations between "n" featured datasets.

In [ ]:

# Importing:
import numpy as np
import pandas as pd
from sklearn.experimental import enable_iterative_imputer  # noqa
from sklearn.impute import IterativeImputer
from collections import defaultdict, Counter
import os
import zipfile
from scipy.stats import mannwhitneyu, chi2_contingency

# Random generator for multiple seeds in the imputation:
RNG = np.random.default_rng()

# Importing: import numpy as np import pandas as pd from sklearn.experimental import enable_iterative_imputer # noqa from sklearn.impute import IterativeImputer from collections import defaultdict, Counter import os import zipfile from scipy.stats import mannwhitneyu, chi2_contingency # Random generator for multiple seeds in the imputation: RNG = np.random.default_rng()

In [ ]:

# Complete main function and auxiliar functions:

# Calculates median and IQR for continuous variables:
def _calculate_stats_continuous(series):
    if series.count() == 0:
        return "N/A (N/A, N/A)"
    med = series.median()
    q1 = series.quantile(0.25)
    q3 = series.quantile(0.75)
    try:
        return f"{med:.2f} ({q1:.2f}, {q3:.2f})"
    except (ValueError, TypeError):
        return "N/A (N/A, N/A)"

# Counts positives registers and % for categorical and binary variables:
def _calculate_stats_binary(series):
    total_non_null = series.count()
    if total_non_null == 0:
        return "0 (0.00%)"
    positive_count = (pd.to_numeric(series, errors='coerce') == 1).sum()
    pct = (positive_count / total_non_null) * 100 if total_non_null > 0 else 0
    return f"{positive_count} ({pct:.2f}%)"

# Calculate P-Value for continuous variables:
def _calculate_p_value_continuous(series_before, series_after):
    s_before = series_before.dropna()
    s_after = series_after.dropna()
    if len(s_before) == 0 or len(s_after) == 0:
        return np.nan
    try:
        stat, p = mannwhitneyu(s_before, s_after, alternative='two-sided')
        return p
    except ValueError:
        return 1.0 if np.array_equal(s_before.values, s_after.values) else np.nan

# Calculate P-Value for categorical and binary variables:
def _calculate_p_value_categorical(series_before, series_after):
    s_before = series_before.dropna()
    s_after = series_after.dropna()
    if len(s_before) == 0 or len(s_after) == 0:
        return np.nan

    # Contingency table counting for all classes:
    counts_before = s_before.value_counts()
    counts_after = s_after.value_counts()
    all_categories = sorted(list(set(counts_before.index).union(set(counts_after.index))))

    # Avoid null categories:
    all_categories = [cat for cat in all_categories if pd.notna(cat)]

    if not all_categories:
        return np.nan

    table = []
    row_before = [counts_before.get(cat, 0) for cat in all_categories]
    row_after = [counts_after.get(cat, 0) for cat in all_categories]
    table = [row_before, row_after]

    if np.sum(table) == 0:
        return np.nan

    try:
        chi2, p, dof, expected = chi2_contingency(table)
        return p
    except ValueError:
        return np.nan

# Detects columns that appear to be dummies generated by pd.get_dummies with a prefix_sep separator.
# Returns a dict: {original_var: [dummy_col1, dummy_col2, ...], ...} and a list of binary_single_cols.
def _get_categorical_blocks(df, prefix_sep='!'):
    blocks = defaultdict(list)
    single_binary = []
    for col in df.columns:
        if prefix_sep in col:
            var, cat = col.split(prefix_sep, 1)
            blocks[var].append(col)
        else:
            unique_nonnull = pd.Series(df[col].dropna().unique())
            if len(unique_nonnull) > 0 and all(x in {0, 1, 0.0, 1.0} for x in unique_nonnull):
                 single_binary.append(col)
            elif pd.api.types.is_numeric_dtype(df[col]):
                 pass
            elif len(unique_nonnull) > 0 and set(unique_nonnull).issubset({0,1}):
                 single_binary.append(col)
    return dict(blocks), single_binary

# Given a 1D array with values ​​(probabilities) for each dummy col in order, decides the final class according to cutoff and tie_strategy.
# Returns selected index and value list after thresholding.
def _reconstruct_categorical_from_dummies(df_row_values, dummy_cols, cutoff, tie_strategy, tie_force, var_name):
    vals = np.array(df_row_values, dtype=float)
    # Apply thresholding to get boolean decisions (for force_binary context)
    ge_cutoff = vals >= cutoff
    num_ge = ge_cutoff.sum()
    if num_ge == 1:
        # unique winner
        winner_idx = int(np.where(ge_cutoff)[0][0])
        return winner_idx, vals
    elif num_ge > 1:
        # tie among >= cutoff
        if tie_strategy == 'highest_probability':
            winner_idx = int(np.argmax(vals))
            return winner_idx, vals
        elif tie_strategy == 'first':
            winner_idx = int(np.where(ge_cutoff)[0][0])
            return winner_idx, vals
        elif tie_strategy == 'nan':
            return None, vals
        elif tie_strategy == 'force_class':
            if tie_force is None:
                return None, vals
            # tie_force may be dict or str: if dict, pick mapping for var_name
            forced = tie_force[var_name] if isinstance(tie_force, dict) else tie_force
            # find index
            forced_idx = None
            for idx, col in enumerate(dummy_cols):
                # dummy col name structure: f"{var}{prefix_sep}{category}"
                # extract category as part after separator (we won't rely on it here; caller may map)
                if forced in col:
                    forced_idx = idx
                    break
            if forced_idx is None:
                # forced class not found -> return None
                return None, vals
            return forced_idx, vals
        else:
            return None, vals
    else:
        # none >= cutoff
        if tie_strategy == 'highest_probability':
            winner_idx = int(np.argmax(vals))
            return winner_idx, vals
        elif tie_strategy == 'first':
            winner_idx = 0
            return winner_idx, vals
        elif tie_strategy == 'nan':
            return None, vals
        elif tie_strategy == 'force_class':
            if tie_force is None:
                return None, vals
            forced = tie_force[var_name] if isinstance(tie_force, dict) else tie_force
            forced_idx = None
            for idx, col in enumerate(dummy_cols):
                if forced in col:
                    forced_idx = idx
                    break
            if forced_idx is None:
                return None, vals
            return forced_idx, vals
        else:
            return None, vals

# Main function:
def MICE_Imputer(
    df,
    interest_vars=None,
    show_results=True,
    random_state=None,
    n=5,
    max_iter=10,
    initial_strategy='most_frequent',
    prefix_sep='!',
    force_binary=False,
    cutoff=0.5,
    cutoff_map=None,
    tie_strategy='highest_probability',
    tie_force=None,
    imputer_kwargs_extra=None
):

    # Validate inputs:
    if interest_vars is None:
        interest_vars = list(df.columns)
    else:
        # Ensure variables exist:
        missing = [v for v in interest_vars if v not in df.columns and not any(v + prefix_sep in c for c in df.columns)]
        if missing:
            raise ValueError(f"Interest variables not found in dataframe columns: {missing}")
    if cutoff_map is None:
        cutoff_map = {}
    if imputer_kwargs_extra is None:
        imputer_kwargs_extra = {}

    # Prepare RNG and per-imputation seeds:
    if random_state is None:
        base_seed = int(RNG.integers(0,2**31-1))
    else:
        base_seed = int(random_state)
    # Create n unique seeds based on base_seed:
    rng_local = np.random.default_rng(base_seed)
    seeds = list(rng_local.integers(0, 2**31-1, size=n))

    # Create a copy of df to avoid altering original:
    df_orig = df.copy(deep=True)
    # Identify original missing mask and counts:
    missing_mask = df_orig.isna()

    # Preprocess: if there are categorical (object / category) columns, convert to dummies:
    df_for_impute = df_orig.copy(deep=True)
    categorical_cols = df_for_impute.select_dtypes(include=['object', 'category']).columns.tolist()
    # For columns that are numeric but contain strings representing categories, they would be object dtype already.
    # Create mapping of original categorical levels for later reconstruction:
    cat_levels = {}
    # We'll produce an expanded dataframe with dummies for categorical columns:
    if len(categorical_cols) > 0:
        for col in categorical_cols:
            # get dummies with prefix and prefix_sep:
            dummies = pd.get_dummies(df_for_impute[col], prefix=col, prefix_sep=prefix_sep, dummy_na=False)
            cat_levels[col] = [c for c in dummies.columns]
            df_for_impute = pd.concat([df_for_impute.drop(columns=[col]), dummies], axis=1)
    else:
        cat_levels = {}

    # After dummying, detect any existing dummy blocks that follow prefix_sep pattern (useful if user already had dummies):
    detected_blocks, single_binary = _get_categorical_blocks(df_for_impute, prefix_sep=prefix_sep)
    # Merge detected_blocks with cat_levels (ensuring union):
    for k, v in detected_blocks.items():
        if k in cat_levels:
            # Extend but avoid duplicates:
            for col in v:
                if col not in cat_levels[k]:
                    cat_levels[k].append(col)
        else:
            cat_levels[k] = v

    # Also consider single_binary columns as categorical single columns (if not already)
    # We'll treat single_binary as categorical blocks of length 1:
    for col in single_binary:
        if col not in df_for_impute.columns:
            continue
        # Only add if not part of existing cat_levels:
        added = False
        for k,v in cat_levels.items():
            if col in v:
                added = True
                break
        if not added:
            cat_levels[col] = [col]

    # Now df_for_impute should be numeric for IterativeImputer (except maybe object leftovers)
    # Try to coerce any remaining non-numeric to numeric when possible (errors='coerce' will produce NaN):
    for col in df_for_impute.columns:
        if not pd.api.types.is_numeric_dtype(df_for_impute[col]):
            df_for_impute[col] = pd.to_numeric(df_for_impute[col], errors='coerce')

    # Start imputations:
    imputed_datasets = []
    stacked_arrays = []
    for i, seed in enumerate(seeds):
        imputer_kwargs = dict(max_iter=max_iter, initial_strategy=initial_strategy, random_state=int(seed))
        # Merge extra kwargs:
        imputer_kwargs.update(imputer_kwargs_extra)
        imputer = IterativeImputer(**imputer_kwargs)
        arr = imputer.fit_transform(df_for_impute.values)
        imputed_array = arr.copy()  # shape (rows, cols)
        stacked_arrays.append(imputed_array)
        # Rebuild DataFrame with same columns as df_for_impute:
        imputed_df_expanded = pd.DataFrame(imputed_array, index=df_for_impute.index, columns=df_for_impute.columns)
        # Now reconstruct categorical variables back into single columns where applicable (cat_levels):
        imputed_df_reconstructed = imputed_df_expanded.copy(deep=True)
        for var, dummy_cols in cat_levels.items():
            # Some entries in dummy_cols might not be present (if some categories absent); filter:
            present = [c for c in dummy_cols if c in imputed_df_expanded.columns]
            if len(present) == 0:
                continue
            # For single-column binary categorical (len==1 and probably 0/1), handle separately:
            if len(present) == 1:
                col_name = present[0]
                # If force_binary, threshold according to cutoff_map or cutoff:
                this_cutoff = cutoff_map.get(var, cutoff)
                if force_binary:
                    binvals = (imputed_df_expanded[col_name] >= this_cutoff).astype('Int64')
                    imputed_df_reconstructed[var] = binvals
                else:
                    # Round to nearest integer where appropriate:
                    imputed_df_reconstructed[var] = imputed_df_expanded[col_name].round().astype('Int64')
                # Drop the dummy column (only if it's not the var name itself):
                if col_name != var:
                    imputed_df_reconstructed = imputed_df_reconstructed.drop(columns=present)
            else:
                # For multi-class blocks: choose class per row based on thresholds and tie_strategy:
                this_cutoff = cutoff_map.get(var, cutoff)
                # Prepare list of categories extracted from dummy column names:
                categories = [c.split(prefix_sep,1)[1] if prefix_sep in c else c for c in present]
                selected = []
                for idx in range(imputed_df_expanded.shape[0]):
                    row_vals = imputed_df_expanded.loc[imputed_df_expanded.index[idx], present].values
                    winner_idx, _ = _reconstruct_categorical_from_dummies(row_vals, present, this_cutoff, tie_strategy, tie_force, var)
                    if winner_idx is None:
                        selected.append(np.nan)
                    else:
                        selected.append(categories[winner_idx])
                imputed_df_reconstructed[var] = pd.Series(selected, index=imputed_df_expanded.index, dtype='object')
                # Drop dummy columns:
                imputed_df_reconstructed = imputed_df_reconstructed.drop(columns=present)
        # Final step: ensure column order similar to original df (put reconstructed categorical columns and numeric columns):
        final_cols = []
        for c in df.columns:
            if c in imputed_df_reconstructed.columns:
                final_cols.append(c)
        # Also append any new numeric columns that were not in original (shouldn't normally happen):
        for c in imputed_df_reconstructed.columns:
            if c not in final_cols:
                final_cols.append(c)
        imputed_df_final = imputed_df_reconstructed[final_cols].copy(deep=True)
        imputed_datasets.append(imputed_df_final)

    # Save imputed datasets:
    zip_filename = 'imputed_datasets.zip'
    print(f"Saving {n} imputed datasets in '{zip_filename}'...")
    try:
        with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zf:
            for i, imp_df in enumerate(imputed_datasets):
                # Save and add to zip file:
                temp_filename = f'imputed_dataset_{i+1}.xlsx'
                imp_df.to_excel(temp_filename, index=True)
                zf.write(temp_filename)
                os.remove(temp_filename)
        print(f"Success: '{zip_filename}' saved.")
    except Exception as e:
        print(f"Error saving zip file: {e}")

    # Stack arrays for pooling (using expanded columns df_for_impute.columns):
    stacked = np.stack(stacked_arrays, axis=0)
    pooled_array = np.mean(stacked, axis=0)

    # Build pooled expanded dataframe (with dummy columns):
    pooled_expanded_df = pd.DataFrame(pooled_array, index=df_for_impute.index, columns=df_for_impute.columns)
    # Reconstruct pooled categorical columns similarly to per-imputation reconstruction:
    pooled_reconstructed = pooled_expanded_df.copy(deep=True)
    for var, dummy_cols in cat_levels.items():
        present = [c for c in dummy_cols if c in pooled_expanded_df.columns]
        if len(present) == 0:
            continue
        if len(present) == 1:
            col_name = present[0]
            this_cutoff = cutoff_map.get(var, cutoff)
            if force_binary:
                pooled_reconstructed[var] = (pooled_expanded_df[col_name] >= this_cutoff).astype('Int64')
            else:
                pooled_reconstructed[var] = pooled_expanded_df[col_name].round().astype('Int64')
            if col_name != var:
                pooled_reconstructed = pooled_reconstructed.drop(columns=present)
        else:
            this_cutoff = cutoff_map.get(var, cutoff)
            categories = [c.split(prefix_sep,1)[1] if prefix_sep in c else c for c in present]
            selected = []
            for idx in range(pooled_expanded_df.shape[0]):
                row_vals = pooled_expanded_df.loc[pooled_expanded_df.index[idx], present].values
                winner_idx, _ = _reconstruct_categorical_from_dummies(row_vals, present, this_cutoff, tie_strategy, tie_force, var)
                if winner_idx is None:
                    selected.append(np.nan)
                else:
                    selected.append(categories[winner_idx])
            pooled_reconstructed[var] = pd.Series(selected, index=pooled_expanded_df.index, dtype='object')
            pooled_reconstructed = pooled_reconstructed.drop(columns=present)

    # Reorder pooled_reconstructed columns to match original df columns where possible:
    pooled_final_cols = [c for c in df.columns if c in pooled_reconstructed.columns]
    for c in pooled_reconstructed.columns:
        if c not in pooled_final_cols:
            pooled_final_cols.append(c)
    pooled_df = pooled_reconstructed[pooled_final_cols].copy(deep=True)

    # Stat analysis:
    stats = []
    for var in interest_vars:
        if var not in df.columns:
            # Skip unknown variables (may happen if interest_vars includes removed dummies):
            continue

        # 1. Count missings (original dataframe):
        count_missing = int(missing_mask[var].sum()) if var in missing_mask.columns else 0
        pct_missing = float(count_missing / len(df) * 100)

        # Skip variable if it have no missing data:
        if count_missing == 0:
            continue

        # 2. Determinate type and calculate metrics (before vs. after imputation):
        is_categorical = var in cat_levels
        series_before = df_orig[var]
        series_after = pooled_df[var]

        stats_before_str = "N/A"
        stats_after_str = "N/A"
        p_value = np.nan
        mean_imputation_variance = np.nan

        if not is_categorical:
            stats_before_str = _calculate_stats_continuous(series_before)
            stats_after_str = _calculate_stats_continuous(series_after)
            p_value = _calculate_p_value_continuous(series_before, series_after)

            # Calculate imputation variance between the multiple datasets:
            if var in df_for_impute.columns:
                col_idx = list(df_for_impute.columns).index(var)
                if count_missing > 0:
                    vals_across = stacked[:, missing_mask[var].values, col_idx]
                    var_per_pos = np.var(vals_across, axis=0, ddof=1)
                    mean_imputation_variance = float(np.nanmean(var_per_pos))
                else:
                    mean_imputation_variance = 0.0 # If there's no missing, variance = 0

        # For categorical and binary variables:
        else:
            stats_before_str = _calculate_stats_binary(series_before)
            stats_after_str = _calculate_stats_binary(series_after)
            p_value = _calculate_p_value_categorical(series_before, series_after)

            # Calculate imputation variance between the multiple datasets:
            dummy_cols = [c for c in cat_levels.get(var, []) if c in df_for_impute.columns]
            if count_missing > 0 and len(dummy_cols) > 0:
                col_indices = [list(df_for_impute.columns).index(c) for c in dummy_cols]
                var_per_pos_list = []
                # Iterate over positions that were originally missing:
                for pos_idx, row_bool in enumerate(missing_mask[var].values):
                    if not row_bool:
                        continue
                    vals = stacked[:, pos_idx, :][:, col_indices]
                    var_per_dummy = np.var(vals, axis=0, ddof=1)
                    var_per_pos_list.append(np.nanmean(var_per_dummy))

                if len(var_per_pos_list) > 0:
                    mean_imputation_variance = float(np.nanmean(var_per_pos_list))
            elif count_missing == 0:
                mean_imputation_variance = 0.0

        # 3. Add stats table:
        stats.append({
            'VARIABLE': var,
            'MISSINGS': f"{count_missing} ({pct_missing:.2f}%)",
            'BEFORE IMPUTATION': stats_before_str,
            'AFTER IMPUTATION': stats_after_str,
            'P-VALUE': p_value,
            'IMPUTATION VARIANCE': mean_imputation_variance
        })

    # Create stat analysis dataframe:
    stats_df = pd.DataFrame(stats)

    if show_results:
        # Show stat analysis table:
        print("Statistical analysis of data imputation with MICE method to numerical and categorical variables")
        # Avoid truncation in terminal:
        with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.width', 1000):\
            display(stats_df)
        print("\nALERT: The table will only show a variable if there is some missing register in its column. Otherwise, it will not be shown in the table.")

    # Return zip file with "n" imputed dataframes, pooled dataframe and stat analysis dataframe:
    return imputed_datasets, pooled_df, stats_df

# Complete main function and auxiliar functions: # Calculates median and IQR for continuous variables: def _calculate_stats_continuous(series): if series.count() == 0: return "N/A (N/A, N/A)" med = series.median() q1 = series.quantile(0.25) q3 = series.quantile(0.75) try: return f"{med:.2f} ({q1:.2f}, {q3:.2f})" except (ValueError, TypeError): return "N/A (N/A, N/A)" # Counts positives registers and % for categorical and binary variables: def _calculate_stats_binary(series): total_non_null = series.count() if total_non_null == 0: return "0 (0.00%)" positive_count = (pd.to_numeric(series, errors='coerce') == 1).sum() pct = (positive_count / total_non_null) * 100 if total_non_null > 0 else 0 return f"{positive_count} ({pct:.2f}%)" # Calculate P-Value for continuous variables: def _calculate_p_value_continuous(series_before, series_after): s_before = series_before.dropna() s_after = series_after.dropna() if len(s_before) == 0 or len(s_after) == 0: return np.nan try: stat, p = mannwhitneyu(s_before, s_after, alternative='two-sided') return p except ValueError: return 1.0 if np.array_equal(s_before.values, s_after.values) else np.nan # Calculate P-Value for categorical and binary variables: def _calculate_p_value_categorical(series_before, series_after): s_before = series_before.dropna() s_after = series_after.dropna() if len(s_before) == 0 or len(s_after) == 0: return np.nan # Contingency table counting for all classes: counts_before = s_before.value_counts() counts_after = s_after.value_counts() all_categories = sorted(list(set(counts_before.index).union(set(counts_after.index)))) # Avoid null categories: all_categories = [cat for cat in all_categories if pd.notna(cat)] if not all_categories: return np.nan table = [] row_before = [counts_before.get(cat, 0) for cat in all_categories] row_after = [counts_after.get(cat, 0) for cat in all_categories] table = [row_before, row_after] if np.sum(table) == 0: return np.nan try: chi2, p, dof, expected = chi2_contingency(table) return p except ValueError: return np.nan # Detects columns that appear to be dummies generated by pd.get_dummies with a prefix_sep separator. # Returns a dict: {original_var: [dummy_col1, dummy_col2, ...], ...} and a list of binary_single_cols. def _get_categorical_blocks(df, prefix_sep='!'): blocks = defaultdict(list) single_binary = [] for col in df.columns: if prefix_sep in col: var, cat = col.split(prefix_sep, 1) blocks[var].append(col) else: unique_nonnull = pd.Series(df[col].dropna().unique()) if len(unique_nonnull) > 0 and all(x in {0, 1, 0.0, 1.0} for x in unique_nonnull): single_binary.append(col) elif pd.api.types.is_numeric_dtype(df[col]): pass elif len(unique_nonnull) > 0 and set(unique_nonnull).issubset({0,1}): single_binary.append(col) return dict(blocks), single_binary # Given a 1D array with values ​​(probabilities) for each dummy col in order, decides the final class according to cutoff and tie_strategy. # Returns selected index and value list after thresholding. def _reconstruct_categorical_from_dummies(df_row_values, dummy_cols, cutoff, tie_strategy, tie_force, var_name): vals = np.array(df_row_values, dtype=float) # Apply thresholding to get boolean decisions (for force_binary context) ge_cutoff = vals >= cutoff num_ge = ge_cutoff.sum() if num_ge == 1: # unique winner winner_idx = int(np.where(ge_cutoff)[0][0]) return winner_idx, vals elif num_ge > 1: # tie among >= cutoff if tie_strategy == 'highest_probability': winner_idx = int(np.argmax(vals)) return winner_idx, vals elif tie_strategy == 'first': winner_idx = int(np.where(ge_cutoff)[0][0]) return winner_idx, vals elif tie_strategy == 'nan': return None, vals elif tie_strategy == 'force_class': if tie_force is None: return None, vals # tie_force may be dict or str: if dict, pick mapping for var_name forced = tie_force[var_name] if isinstance(tie_force, dict) else tie_force # find index forced_idx = None for idx, col in enumerate(dummy_cols): # dummy col name structure: f"{var}{prefix_sep}{category}" # extract category as part after separator (we won't rely on it here; caller may map) if forced in col: forced_idx = idx break if forced_idx is None: # forced class not found -> return None return None, vals return forced_idx, vals else: return None, vals else: # none >= cutoff if tie_strategy == 'highest_probability': winner_idx = int(np.argmax(vals)) return winner_idx, vals elif tie_strategy == 'first': winner_idx = 0 return winner_idx, vals elif tie_strategy == 'nan': return None, vals elif tie_strategy == 'force_class': if tie_force is None: return None, vals forced = tie_force[var_name] if isinstance(tie_force, dict) else tie_force forced_idx = None for idx, col in enumerate(dummy_cols): if forced in col: forced_idx = idx break if forced_idx is None: return None, vals return forced_idx, vals else: return None, vals # Main function: def MICE_Imputer( df, interest_vars=None, show_results=True, random_state=None, n=5, max_iter=10, initial_strategy='most_frequent', prefix_sep='!', force_binary=False, cutoff=0.5, cutoff_map=None, tie_strategy='highest_probability', tie_force=None, imputer_kwargs_extra=None ): # Validate inputs: if interest_vars is None: interest_vars = list(df.columns) else: # Ensure variables exist: missing = [v for v in interest_vars if v not in df.columns and not any(v + prefix_sep in c for c in df.columns)] if missing: raise ValueError(f"Interest variables not found in dataframe columns: {missing}") if cutoff_map is None: cutoff_map = {} if imputer_kwargs_extra is None: imputer_kwargs_extra = {} # Prepare RNG and per-imputation seeds: if random_state is None: base_seed = int(RNG.integers(0,2**31-1)) else: base_seed = int(random_state) # Create n unique seeds based on base_seed: rng_local = np.random.default_rng(base_seed) seeds = list(rng_local.integers(0, 2**31-1, size=n)) # Create a copy of df to avoid altering original: df_orig = df.copy(deep=True) # Identify original missing mask and counts: missing_mask = df_orig.isna() # Preprocess: if there are categorical (object / category) columns, convert to dummies: df_for_impute = df_orig.copy(deep=True) categorical_cols = df_for_impute.select_dtypes(include=['object', 'category']).columns.tolist() # For columns that are numeric but contain strings representing categories, they would be object dtype already. # Create mapping of original categorical levels for later reconstruction: cat_levels = {} # We'll produce an expanded dataframe with dummies for categorical columns: if len(categorical_cols) > 0: for col in categorical_cols: # get dummies with prefix and prefix_sep: dummies = pd.get_dummies(df_for_impute[col], prefix=col, prefix_sep=prefix_sep, dummy_na=False) cat_levels[col] = [c for c in dummies.columns] df_for_impute = pd.concat([df_for_impute.drop(columns=[col]), dummies], axis=1) else: cat_levels = {} # After dummying, detect any existing dummy blocks that follow prefix_sep pattern (useful if user already had dummies): detected_blocks, single_binary = _get_categorical_blocks(df_for_impute, prefix_sep=prefix_sep) # Merge detected_blocks with cat_levels (ensuring union): for k, v in detected_blocks.items(): if k in cat_levels: # Extend but avoid duplicates: for col in v: if col not in cat_levels[k]: cat_levels[k].append(col) else: cat_levels[k] = v # Also consider single_binary columns as categorical single columns (if not already) # We'll treat single_binary as categorical blocks of length 1: for col in single_binary: if col not in df_for_impute.columns: continue # Only add if not part of existing cat_levels: added = False for k,v in cat_levels.items(): if col in v: added = True break if not added: cat_levels[col] = [col] # Now df_for_impute should be numeric for IterativeImputer (except maybe object leftovers) # Try to coerce any remaining non-numeric to numeric when possible (errors='coerce' will produce NaN): for col in df_for_impute.columns: if not pd.api.types.is_numeric_dtype(df_for_impute[col]): df_for_impute[col] = pd.to_numeric(df_for_impute[col], errors='coerce') # Start imputations: imputed_datasets = [] stacked_arrays = [] for i, seed in enumerate(seeds): imputer_kwargs = dict(max_iter=max_iter, initial_strategy=initial_strategy, random_state=int(seed)) # Merge extra kwargs: imputer_kwargs.update(imputer_kwargs_extra) imputer = IterativeImputer(**imputer_kwargs) arr = imputer.fit_transform(df_for_impute.values) imputed_array = arr.copy() # shape (rows, cols) stacked_arrays.append(imputed_array) # Rebuild DataFrame with same columns as df_for_impute: imputed_df_expanded = pd.DataFrame(imputed_array, index=df_for_impute.index, columns=df_for_impute.columns) # Now reconstruct categorical variables back into single columns where applicable (cat_levels): imputed_df_reconstructed = imputed_df_expanded.copy(deep=True) for var, dummy_cols in cat_levels.items(): # Some entries in dummy_cols might not be present (if some categories absent); filter: present = [c for c in dummy_cols if c in imputed_df_expanded.columns] if len(present) == 0: continue # For single-column binary categorical (len==1 and probably 0/1), handle separately: if len(present) == 1: col_name = present[0] # If force_binary, threshold according to cutoff_map or cutoff: this_cutoff = cutoff_map.get(var, cutoff) if force_binary: binvals = (imputed_df_expanded[col_name] >= this_cutoff).astype('Int64') imputed_df_reconstructed[var] = binvals else: # Round to nearest integer where appropriate: imputed_df_reconstructed[var] = imputed_df_expanded[col_name].round().astype('Int64') # Drop the dummy column (only if it's not the var name itself): if col_name != var: imputed_df_reconstructed = imputed_df_reconstructed.drop(columns=present) else: # For multi-class blocks: choose class per row based on thresholds and tie_strategy: this_cutoff = cutoff_map.get(var, cutoff) # Prepare list of categories extracted from dummy column names: categories = [c.split(prefix_sep,1)[1] if prefix_sep in c else c for c in present] selected = [] for idx in range(imputed_df_expanded.shape[0]): row_vals = imputed_df_expanded.loc[imputed_df_expanded.index[idx], present].values winner_idx, _ = _reconstruct_categorical_from_dummies(row_vals, present, this_cutoff, tie_strategy, tie_force, var) if winner_idx is None: selected.append(np.nan) else: selected.append(categories[winner_idx]) imputed_df_reconstructed[var] = pd.Series(selected, index=imputed_df_expanded.index, dtype='object') # Drop dummy columns: imputed_df_reconstructed = imputed_df_reconstructed.drop(columns=present) # Final step: ensure column order similar to original df (put reconstructed categorical columns and numeric columns): final_cols = [] for c in df.columns: if c in imputed_df_reconstructed.columns: final_cols.append(c) # Also append any new numeric columns that were not in original (shouldn't normally happen): for c in imputed_df_reconstructed.columns: if c not in final_cols: final_cols.append(c) imputed_df_final = imputed_df_reconstructed[final_cols].copy(deep=True) imputed_datasets.append(imputed_df_final) # Save imputed datasets: zip_filename = 'imputed_datasets.zip' print(f"Saving {n} imputed datasets in '{zip_filename}'...") try: with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zf: for i, imp_df in enumerate(imputed_datasets): # Save and add to zip file: temp_filename = f'imputed_dataset_{i+1}.xlsx' imp_df.to_excel(temp_filename, index=True) zf.write(temp_filename) os.remove(temp_filename) print(f"Success: '{zip_filename}' saved.") except Exception as e: print(f"Error saving zip file: {e}") # Stack arrays for pooling (using expanded columns df_for_impute.columns): stacked = np.stack(stacked_arrays, axis=0) pooled_array = np.mean(stacked, axis=0) # Build pooled expanded dataframe (with dummy columns): pooled_expanded_df = pd.DataFrame(pooled_array, index=df_for_impute.index, columns=df_for_impute.columns) # Reconstruct pooled categorical columns similarly to per-imputation reconstruction: pooled_reconstructed = pooled_expanded_df.copy(deep=True) for var, dummy_cols in cat_levels.items(): present = [c for c in dummy_cols if c in pooled_expanded_df.columns] if len(present) == 0: continue if len(present) == 1: col_name = present[0] this_cutoff = cutoff_map.get(var, cutoff) if force_binary: pooled_reconstructed[var] = (pooled_expanded_df[col_name] >= this_cutoff).astype('Int64') else: pooled_reconstructed[var] = pooled_expanded_df[col_name].round().astype('Int64') if col_name != var: pooled_reconstructed = pooled_reconstructed.drop(columns=present) else: this_cutoff = cutoff_map.get(var, cutoff) categories = [c.split(prefix_sep,1)[1] if prefix_sep in c else c for c in present] selected = [] for idx in range(pooled_expanded_df.shape[0]): row_vals = pooled_expanded_df.loc[pooled_expanded_df.index[idx], present].values winner_idx, _ = _reconstruct_categorical_from_dummies(row_vals, present, this_cutoff, tie_strategy, tie_force, var) if winner_idx is None: selected.append(np.nan) else: selected.append(categories[winner_idx]) pooled_reconstructed[var] = pd.Series(selected, index=pooled_expanded_df.index, dtype='object') pooled_reconstructed = pooled_reconstructed.drop(columns=present) # Reorder pooled_reconstructed columns to match original df columns where possible: pooled_final_cols = [c for c in df.columns if c in pooled_reconstructed.columns] for c in pooled_reconstructed.columns: if c not in pooled_final_cols: pooled_final_cols.append(c) pooled_df = pooled_reconstructed[pooled_final_cols].copy(deep=True) # Stat analysis: stats = [] for var in interest_vars: if var not in df.columns: # Skip unknown variables (may happen if interest_vars includes removed dummies): continue # 1. Count missings (original dataframe): count_missing = int(missing_mask[var].sum()) if var in missing_mask.columns else 0 pct_missing = float(count_missing / len(df) * 100) # Skip variable if it have no missing data: if count_missing == 0: continue # 2. Determinate type and calculate metrics (before vs. after imputation): is_categorical = var in cat_levels series_before = df_orig[var] series_after = pooled_df[var] stats_before_str = "N/A" stats_after_str = "N/A" p_value = np.nan mean_imputation_variance = np.nan if not is_categorical: stats_before_str = _calculate_stats_continuous(series_before) stats_after_str = _calculate_stats_continuous(series_after) p_value = _calculate_p_value_continuous(series_before, series_after) # Calculate imputation variance between the multiple datasets: if var in df_for_impute.columns: col_idx = list(df_for_impute.columns).index(var) if count_missing > 0: vals_across = stacked[:, missing_mask[var].values, col_idx] var_per_pos = np.var(vals_across, axis=0, ddof=1) mean_imputation_variance = float(np.nanmean(var_per_pos)) else: mean_imputation_variance = 0.0 # If there's no missing, variance = 0 # For categorical and binary variables: else: stats_before_str = _calculate_stats_binary(series_before) stats_after_str = _calculate_stats_binary(series_after) p_value = _calculate_p_value_categorical(series_before, series_after) # Calculate imputation variance between the multiple datasets: dummy_cols = [c for c in cat_levels.get(var, []) if c in df_for_impute.columns] if count_missing > 0 and len(dummy_cols) > 0: col_indices = [list(df_for_impute.columns).index(c) for c in dummy_cols] var_per_pos_list = [] # Iterate over positions that were originally missing: for pos_idx, row_bool in enumerate(missing_mask[var].values): if not row_bool: continue vals = stacked[:, pos_idx, :][:, col_indices] var_per_dummy = np.var(vals, axis=0, ddof=1) var_per_pos_list.append(np.nanmean(var_per_dummy)) if len(var_per_pos_list) > 0: mean_imputation_variance = float(np.nanmean(var_per_pos_list)) elif count_missing == 0: mean_imputation_variance = 0.0 # 3. Add stats table: stats.append({ 'VARIABLE': var, 'MISSINGS': f"{count_missing} ({pct_missing:.2f}%)", 'BEFORE IMPUTATION': stats_before_str, 'AFTER IMPUTATION': stats_after_str, 'P-VALUE': p_value, 'IMPUTATION VARIANCE': mean_imputation_variance }) # Create stat analysis dataframe: stats_df = pd.DataFrame(stats) if show_results: # Show stat analysis table: print("Statistical analysis of data imputation with MICE method to numerical and categorical variables") # Avoid truncation in terminal: with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.width', 1000):\ display(stats_df) print("\nALERT: The table will only show a variable if there is some missing register in its column. Otherwise, it will not be shown in the table.") # Return zip file with "n" imputed dataframes, pooled dataframe and stat analysis dataframe: return imputed_datasets, pooled_df, stats_df

In [ ]:

# Example:

path_age = 'age_weighted_missing_data_sample_signs_comorb_2019.xlsx'
path_random = 'random_missing_data_sample_signs_comorb_2019.xlsx'

df_age = pd.read_excel(path_age)
df_random = pd.read_excel(path_random)

print("--- Processing df_age ---")
imputed_age_list, pooled_age, stats_age = MICE_Imputer(df_age, n=5, show_results=True)

print("\n--- Processing df_random ---")
imputed_random_list, pooled_random, stats_random = MICE_Imputer(df_random, show_results=True)

print("\nSaving output files (pooled and stats)...")
pooled_age.to_excel('pooled_age.xlsx', index=True)
stats_age.to_excel('stats_age.xlsx', index=True)

pooled_random.to_excel('pooled_random.xlsx', index=True)
stats_random.to_excel('stats_random.xlsx', index=True)

print("\nProccess conclued. Verify files 'imputed_datasets.zip', 'pooled_*.xlsx' and 'stats_*.xlsx' in directory.")

# Example: path_age = 'age_weighted_missing_data_sample_signs_comorb_2019.xlsx' path_random = 'random_missing_data_sample_signs_comorb_2019.xlsx' df_age = pd.read_excel(path_age) df_random = pd.read_excel(path_random) print("--- Processing df_age ---") imputed_age_list, pooled_age, stats_age = MICE_Imputer(df_age, n=5, show_results=True) print("\n--- Processing df_random ---") imputed_random_list, pooled_random, stats_random = MICE_Imputer(df_random, show_results=True) print("\nSaving output files (pooled and stats)...") pooled_age.to_excel('pooled_age.xlsx', index=True) stats_age.to_excel('stats_age.xlsx', index=True) pooled_random.to_excel('pooled_random.xlsx', index=True) stats_random.to_excel('stats_random.xlsx', index=True) print("\nProccess conclued. Verify files 'imputed_datasets.zip', 'pooled_*.xlsx' and 'stats_*.xlsx' in directory.")