applications/offline/mi.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, math, random, time\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import statsmodels.api as sm\n",
    "import tomllib\n",
    "\n",
    "from copy import copy\n",
    "from tqdm import tqdm\n",
    "from sklearn.linear_model import LinearRegression, LogisticRegression\n",
    "from sklearn.experimental import enable_iterative_imputer\n",
    "from sklearn.impute import IterativeImputer\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../config/mi.toml\", \"rb\") as f:\n",
    "    config = tomllib.load(f)\n",
    "\n",
    "run_scenarios = config[\"run_scenarios\"]\n",
    "run_pymice = config[\"run_pymice\"]\n",
    "run_simice = config[\"run_simice\"]\n",
    "run_real_setup = config[\"run_real_setup\"]\n",
    "\n",
    "pymice_runs_count = config[\"pymice_runs_count\"]\n",
    "simice_runs_count = config[\"simice_runs_count\"]\n",
    "real_setup_runs_count = config[\"real_setup_runs_count\"]\n",
    "\n",
    "mi_factor = config[\"mi_factor\"]\n",
    "miss_val = config[\"miss_val\"]\n",
    "dump_data = config[\"dump_data\"]\n",
    "labels_noise_scale = config[\"im_noise_scale\"]\n",
    "\n",
    "paper_scenario_rows = config[\"paper_scenario_rows\"]\n",
    "midn_scenario_rows = config[\"midn_scenario_rows\"]\n",
    "midn_scenario_cols = config[\"midn_scenario_cols\"]\n",
    "midn_miss_rows_count = config[\"midn_miss_rows_count\"]\n",
    "\n",
    "scenarios_runs_count = config[\"scenarios_runs_count\"]\n",
    "scenarios_miss_col = config[\"scenarios_miss_col\"]\n",
    "\n",
    "real_setup_outlier_threshold = config[\"real_setup_outlier_threshold\"]\n",
    "\n",
    "random.seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Data generators/parsers \"\"\"\n",
    "\n",
    "def export_data(name, tampered_data, non_tampered_labels, non_tampered_data, miss_rows, miss_val):\n",
    "    path = f\"../../data/mi/{name}/{len(tampered_data)}\"\n",
    "    Path(path).mkdir(parents=True, exist_ok=True)\n",
    "    Path(\"../../data/mi/dump\").mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "    np.savetxt(f\"{path}/tampered_data.txt\", np.nan_to_num(tampered_data, nan=miss_val))\n",
    "    np.savetxt(f\"{path}/tampered_data_w_nans.txt\", tampered_data)\n",
    "    np.savetxt(f\"{path}/non_tampered_data.txt\", non_tampered_data)\n",
    "    np.savetxt(f\"{path}/non_tampered_labels.txt\", non_tampered_labels)\n",
    "    np.savetxt(f\"{path}/miss_rows.txt\", miss_rows, fmt='%d')\n",
    "\n",
    "def import_data(path):\n",
    "    tampered_data = np.loadtxt(f\"{path}/tampered_data_w_nans.txt\")\n",
    "    non_tampered_data = np.loadtxt(f\"{path}/non_tampered_data.txt\")\n",
    "    non_tampered_labels = np.expand_dims(np.loadtxt(f\"{path}/non_tampered_labels.txt\"), axis=1)\n",
    "    miss_rows = np.loadtxt(f\"{path}/miss_rows.txt\").astype(int).tolist()\n",
    "    \n",
    "    return tampered_data, non_tampered_labels, non_tampered_data, miss_rows\n",
    "\n",
    "def batch(data, batch_size: int = 0, batch_count: int = 0):\n",
    "    assert bool(batch_size) ^ bool(batch_count), \"Invalid batching parameters: either size or count should be passed\"\n",
    "\n",
    "    data_size = data.shape[0]\n",
    "    if batch_size == 0:\n",
    "        batch_size = data_size // batch_count\n",
    "    elif batch_count == 0:\n",
    "        batch_count = (data_size + batch_size - 1) // batch_size\n",
    "    \n",
    "    return [data[i * batch_size:(i + 1) * batch_size] for i in range(batch_count)]\n",
    "\n",
    "def relative_noise(arr, scale):\n",
    "    if scale == 0.0:\n",
    "        return np.zeros_like(arr)\n",
    "    \n",
    "    return np.random.normal(loc=0.0, scale=np.abs(arr) * scale)\n",
    "\n",
    "def load_or_generate_paper_data(rows: int, scenario: str, miss_col: int, labels_noise_scale: float):\n",
    "    # Generate data as in https://www.nature.com/articles/s41467-020-19270-2 scenarios\n",
    "    assert scenario in {\"scenario_1\", \"scenario_2\"}, f\"Invalid scenario for MI data {scenario}\"\n",
    "\n",
    "    feature_count = 2\n",
    "    scenario_name = f\"paper_{scenario}\"\n",
    "    path = f\"../../data/mi/{scenario_name}/{rows}\"\n",
    "    if os.path.isdir(path):\n",
    "        return scenario_name, *import_data(path)\n",
    "\n",
    "    # x_2 is from U(-3, 3)\n",
    "    x_2 = np.random.rand(rows) * 6 - 3\n",
    "    # Normalize to avoid MPC overflow later on\n",
    "    if rows > 500:\n",
    "        x_2 /= float(rows // 500)\n",
    "\n",
    "    if scenario == \"linear\":\n",
    "        # Linear x_1 is from normal(mu=0.2 - 0.5*x_2, sigma=1.0)\n",
    "        x_1 = np.random.normal(loc=0.2 - 0.5 * x_2, scale=1.0)\n",
    "    else:\n",
    "        # Logistic x_1 is from bernoulli(n=1, p=1 / (1 + exp(-0.2 + 0.5*x_2)))\n",
    "        x_1 = np.random.binomial(n=1, p=1 / (1 + np.exp(-0.2 + 0.5 * x_2)))\n",
    "\n",
    "    non_tampered_data = np.hstack((np.expand_dims(x_1, axis=1), np.expand_dims(x_2, axis=1)))\n",
    "    \n",
    "    miss_rows = []\n",
    "    tampered_data = non_tampered_data.copy()\n",
    "    \n",
    "    non_tampered_labels = non_tampered_data @ np.ones((feature_count, 1)) + 1.0  # Weights and bias are ones\n",
    "    non_tampered_labels += relative_noise(non_tampered_labels, scale=labels_noise_scale)\n",
    "    for i in range(rows):\n",
    "        # The variable is missing with probability 1 / (1 + math.exp(-0.3 + 0.2 * labels[i][0] - 0.1 * x_2[i]))\n",
    "        if (1 / (1 + math.exp(-0.3 + 0.2 * non_tampered_labels[i][0] - 0.1 * x_2[i]))) > random.random():\n",
    "            tampered_data[i][miss_col] = np.nan\n",
    "            miss_rows.append(i)\n",
    "\n",
    "    print(f\"Generated data missing rate: {len(miss_rows)}/{len(non_tampered_data)}\")\n",
    "\n",
    "    scenario_tuple = (scenario_name, tampered_data, non_tampered_labels, non_tampered_data, miss_rows)\n",
    "    if dump_data:\n",
    "        export_data(*scenario_tuple, miss_val)\n",
    "\n",
    "    return scenario_tuple\n",
    "\n",
    "def load_real_data(outlier_threshold):\n",
    "    # Real data\n",
    "    df = pd.read_csv(\"../../data/mi/real/gcasr.csv\")\n",
    "    binary_columns = [\"RaceAA\", \"RaceW\", \"HlthInsM\", \"HlthInsN\", \"Day\", \"NPO\", \"MedHisST\", \"MedHisTI\", \"MedHisVP\", \"MedHisFHSTK\"]\n",
    "    continuous_columns = [\"NIHStrkS\", \"EMSNote\", \"LipTotal\", \"Age\", \"Gender\"]\n",
    "    col_types = [\"lin\" if i < len(continuous_columns) else \"log\" for i in range(len(binary_columns) + len(continuous_columns))]\n",
    "    df_of_interest = df[continuous_columns + binary_columns]\n",
    "\n",
    "    # Normalize continuous columns\n",
    "    df_of_interest[continuous_columns] = (df_of_interest[continuous_columns] - df_of_interest[continuous_columns].mean()) / df_of_interest[continuous_columns].std()\n",
    "\n",
    "    # Fill the missing data for numerical feasibility\n",
    "    real_data = df_of_interest.to_numpy()\n",
    "    # Drop the missing data to get the training data\n",
    "    real_complete_data = df_of_interest.dropna()\n",
    "\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/real/incomplete_data.txt\", np.nan_to_num(real_data, nan=miss_val), fmt='%10.2f')\n",
    "        np.savetxt(f\"../../data/mi/real/complete_data.txt\", real_complete_data.to_numpy(), fmt='%10.2f')\n",
    "\n",
    "    non_miss_rows = real_complete_data.index.to_numpy()\n",
    "\n",
    "    # Get labels\n",
    "    x = df[\"x\"].fillna(0)\n",
    "    # Reduce the outliers\n",
    "    x[x < -outlier_threshold] = -outlier_threshold\n",
    "    x[x > outlier_threshold] = outlier_threshold\n",
    "    # Normalize\n",
    "    x = (x - x.mean()) / x.std()\n",
    "    real_labels = x.to_numpy()\n",
    "\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/real/labels.txt\", real_labels, fmt='%10.2f')\n",
    "\n",
    "    # Get missing rows\n",
    "    real_miss_rows = df_of_interest.isna().to_numpy().astype(int).T\n",
    "\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/real/miss_rows.txt\", real_miss_rows, fmt='%d')\n",
    "\n",
    "    return real_data, real_complete_data.to_numpy(), real_labels, real_miss_rows, col_types, non_miss_rows\n",
    "\n",
    "def load_or_generate_midn_data(rows, cols, scenario, miss_rows, miss_col, labels_noise_scale):\n",
    "    assert scenario in {\"scenario_1\", \"scenario_2\"}, f\"Invalid scenario for MI data {scenario}\"\n",
    "    \n",
    "    scenario_name = f\"midn_{scenario}\"\n",
    "    path = f\"../../data/mi/{scenario_name}/{rows}\"\n",
    "    if os.path.isdir(path):\n",
    "        return scenario_name, *import_data(path)\n",
    "\n",
    "    non_tampered_data = np.random.normal(size=(rows, cols))\n",
    "\n",
    "    if \"_2\" or \"_4\" in scenario:\n",
    "        non_tampered_data[:, 0] = np.random.binomial(1, 0.5, rows)\n",
    "\n",
    "    non_tampered_labels = non_tampered_data @ np.ones((cols, 1)) + 1.0  # Weights and bias are ones\n",
    "    non_tampered_labels += relative_noise(non_tampered_labels, scale=labels_noise_scale)\n",
    "\n",
    "    tampered_data = non_tampered_data.copy()\n",
    "    for i in range(miss_rows):\n",
    "        tampered_data[i][miss_col] = np.nan\n",
    "    \n",
    "    scenario_tuple = (scenario_name, tampered_data, non_tampered_labels, non_tampered_data, list(range(miss_rows)))\n",
    "    if dump_data:\n",
    "        export_data(*scenario_tuple, miss_val)\n",
    "    \n",
    "    return scenario_tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Imputers \"\"\"\n",
    "\n",
    "class Imputer:\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "    \n",
    "    @staticmethod\n",
    "    def split_train_test(data, target_col: int, miss_rows):\n",
    "        m, n = data.shape\n",
    "        miss_count = len(miss_rows)\n",
    "\n",
    "        complete_rows = np.zeros((m - miss_count, n - 1))\n",
    "        complete_labels = np.zeros((m - miss_count,))\n",
    "        incomplete_rows = np.zeros((miss_count, n - 1))\n",
    "\n",
    "        x_idx = 0\n",
    "        x_com_idx = 0\n",
    "        for i in range(m):\n",
    "            row = data[i]\n",
    "            row_complement = np.delete(row, target_col)\n",
    "            if i in miss_rows:\n",
    "                incomplete_rows[x_com_idx] = row_complement\n",
    "                x_com_idx += 1\n",
    "            else:\n",
    "                complete_rows[x_idx] = row_complement\n",
    "                complete_labels[x_idx] = row[target_col]\n",
    "                x_idx += 1\n",
    "        \n",
    "        return complete_rows, incomplete_rows, complete_labels\n",
    "    \n",
    "    def fit(self, complete_data, complete_labels):\n",
    "        self.model.fit(X=complete_data, y=complete_labels)\n",
    "    \n",
    "    def predict(self, data):\n",
    "        return self.model.predict_proba(data) if isinstance(self.model, LogisticRegression) else self.model.predict(data)\n",
    "    \n",
    "    def impute(self, data, miss_rows, miss_col: int, noise: bool):\n",
    "        complete_data, incomplete_data, complete_labels = Imputer.split_train_test(data, miss_col, miss_rows)\n",
    "        \n",
    "        if len(incomplete_data) == 0:\n",
    "            return complete_data.copy()\n",
    "        \n",
    "        self.model.fit(X=complete_data, y=complete_labels)\n",
    "        imputed_y = self.predict(incomplete_data)\n",
    "        if noise:\n",
    "            imputed_y += np.random.normal(0.0, 0.01, size=imputed_y.shape)\n",
    "\n",
    "        idx = 0\n",
    "        imputed_data = data.copy()\n",
    "        for i in miss_rows:\n",
    "            imputed_data[i, miss_col] = imputed_y.ravel()[idx]\n",
    "            idx += 1\n",
    "\n",
    "        return imputed_data\n",
    "    \n",
    "    def impute(self, data, incomplete_data, miss_rows, miss_col, noise):\n",
    "        imputed_y = self.predict(incomplete_data)\n",
    "        if noise:\n",
    "            imputed_y += np.random.normal(0.0, 0.01, size=imputed_y.shape)\n",
    "\n",
    "        idx = 0\n",
    "        imputed_data = data.copy()\n",
    "        for i in miss_rows:\n",
    "            imputed_data[i, miss_col] = imputed_y.ravel()[idx]\n",
    "            idx += 1\n",
    "\n",
    "        return imputed_data\n",
    "    \n",
    "    def impute_inplace(self, data, mask, miss_col, noise):\n",
    "        if mask.sum() == 0:  # Nothing to impute --- no missing data\n",
    "            return\n",
    "        \n",
    "        reference_data = np.hstack((data[:, :miss_col], data[:, miss_col+1:]))\n",
    "        imputed_y = self.predict(reference_data[np.where(mask == 1)])\n",
    "        if noise and isinstance(self.model, LinearRegression):\n",
    "            imputed_y += np.random.normal(0.0, 0.01, size=imputed_y.shape)\n",
    "\n",
    "        idx = 0\n",
    "        for i in range(len(mask)):\n",
    "            if mask[i]:\n",
    "                data[i, miss_col] = imputed_y.ravel()[idx]\n",
    "                idx += 1\n",
    "    \n",
    "\n",
    "class MI:\n",
    "    def __init__(self, factor: int, impute_model, fit_model):\n",
    "        self.imputer = Imputer(impute_model)\n",
    "        self.model = fit_model\n",
    "        self.factor = factor\n",
    "\n",
    "    @staticmethod\n",
    "    def rubin(weights):\n",
    "        return np.mean(weights, axis=0)\n",
    "    \n",
    "    def fit(self, data, labels, miss_rows, miss_col, noise):\n",
    "        complete_data, incomplete_data, complete_labels = Imputer.split_train_test(data, miss_col, miss_rows)\n",
    "        \n",
    "        imputed_data = []\n",
    "        self.imputer.fit(complete_data=complete_data, complete_labels=complete_labels)\n",
    "        \n",
    "        imputed_data = [self.imputer.impute(\n",
    "            data=data, incomplete_data=incomplete_data,\n",
    "            miss_rows=miss_rows, miss_col=miss_col, noise=noise) for _ in range(self.factor)]\n",
    "        \n",
    "        self.fit_analysis_model(imputed_data, labels)\n",
    "        return self, imputed_data\n",
    "    \n",
    "    def fit_analysis_model(self, imputed_data, labels):\n",
    "        weights = np.zeros((self.factor, imputed_data[0].shape[1]))\n",
    "        for i, data in enumerate(imputed_data):\n",
    "            self.model.fit(X=data, y=labels)\n",
    "            weights[i] = self.model.coef_.copy()\n",
    "        \n",
    "        self.model.coef_ = MI.rubin(weights)\n",
    "    \n",
    "\n",
    "class BatchMI:\n",
    "    def __init__(self, factor: int, impute_models, fit_model):\n",
    "        self.imputers = [Imputer(impute_model) for impute_model in impute_models]\n",
    "        self.model = fit_model\n",
    "        self.factor = factor\n",
    "\n",
    "    def fit(self, data, complete_data, labels, miss_rows):\n",
    "        miss_col = 0\n",
    "        for imputer in self.imputers:\n",
    "            train_data = np.hstack((complete_data[:, :miss_col], complete_data[:, miss_col+1:]))\n",
    "            train_labels = complete_data[:, miss_col:miss_col+1]\n",
    "            \n",
    "            imputer.fit(complete_data=train_data, complete_labels=train_labels.ravel())\n",
    "            miss_col += 1\n",
    "        \n",
    "        imputed_data = []\n",
    "        weights = np.zeros((self.factor, data.shape[1]))\n",
    "        for i in range(self.factor):\n",
    "            inter_imputed_data = data.copy()\n",
    "            self.impute_inplace(inter_imputed_data, miss_rows)\n",
    "            imputed_data.append(inter_imputed_data)\n",
    "            self.model.fit(X=inter_imputed_data, y=labels)\n",
    "            weights[i] = self.model.coef_.copy()\n",
    "        \n",
    "        self.model.coef_ = MI.rubin(weights)\n",
    "        return self, imputed_data\n",
    "    \n",
    "    def impute_inplace(self, incomplete_data, miss_rows):\n",
    "        miss_col = 0\n",
    "        for imputer in self.imputers:\n",
    "            imputer.impute_inplace(\n",
    "                data=incomplete_data,\n",
    "                mask=miss_rows[miss_col], miss_col=miss_col, noise=True)\n",
    "            miss_col += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Closed-form linear regression w.r.t. least squares\"\"\"\n",
    "\n",
    "class ClosedFormLinReg:\n",
    "    def fit(self, X, y):\n",
    "        X_tilde = np.hstack((X, np.ones((X.shape[0], 1))))\n",
    "        coef = np.linalg.inv(X_tilde.T @ X_tilde) @ X_tilde.T @ y\n",
    "        self.coef_ = coef[:-1].flatten()\n",
    "        self.intercept = coef[-1]\n",
    "\n",
    "    def predict(self, X):\n",
    "        return X @ np.expand_dims(self.coef_, axis=1) + self.intercept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Custom logistic regression \"\"\"\n",
    "\n",
    "def sigmoid(x):\n",
    "    return 1 / (1 + np.exp(-x))\n",
    "\n",
    "\n",
    "class CustomLogReg:\n",
    "    def __init__(self, optimizer: str = \"bgd\"):\n",
    "        self.optimizer = optimizer\n",
    "    \n",
    "    def fit(self, X, y, step: float = 1 / (1 << 7), epochs: int = 10, verbose: bool = False):\n",
    "        self.coef_ = CustomLogReg._fit(X, y, np.random.rand(X.shape[1] + 1, 1), step, epochs, self.optimizer, verbose)\n",
    "        return self\n",
    "    \n",
    "    def predict(self, X):\n",
    "        X_tilde = np.hstack((X, np.ones((X.shape[0], 1))))\n",
    "        return np.round(sigmoid((X_tilde @ self.coef_)))\n",
    "    \n",
    "    @staticmethod\n",
    "    def loss(X, y, w):\n",
    "        return CustomLogReg.log_loss(X, y, w)\n",
    "\n",
    "    def _fit(X, y, initial_w, step: float, epochs: int, optimizer: str, verbose: bool):\n",
    "        # Adding bias\n",
    "        X_tilde = np.hstack((X, np.ones((X.shape[0], 1))))\n",
    "\n",
    "        # Gradient descent\n",
    "        if optimizer == \"bgd\":\n",
    "            return CustomLogReg._bgd(X_tilde, y, initial_w, step, epochs, verbose)\n",
    "        if optimizer == \"mbgd\":\n",
    "            return CustomLogReg._mbgd(X_tilde, y, initial_w, step, epochs, 10, verbose)\n",
    "        else:\n",
    "            raise ValueError(f\"CustomLogReg: invalid optimizer passed: {optimizer}\")\n",
    "    \n",
    "    def _bgd(X_tilde, y, initial_w, step: float, epochs: int, verbose: bool):\n",
    "        if verbose:\n",
    "            print(f\"Log. reg. BGD step size:\", step)\n",
    "        \n",
    "        # Batched gradient descent\n",
    "        w = initial_w  # n x 1\n",
    "        y_mat = np.expand_dims(y.flatten(), axis=1)\n",
    "        for _ in range(epochs):\n",
    "            if verbose:\n",
    "                print(f\"Log. reg. BGD epoch {_ + 1}/{epochs}\")\n",
    "                print(f\"\\t weigts avg {np.mean(w)} | loss: {CustomLogReg.loss(X_tilde, y, w)}\")\n",
    "            \n",
    "            dot = X_tilde @ w  # m x 1\n",
    "            sig = sigmoid(dot)  # m x 1\n",
    "            w -= X_tilde.T @ (sig - y_mat) * step  # n x 1\n",
    "        \n",
    "        return w\n",
    "    \n",
    "    def _mbgd(X_tilde, y, initial_w, step: float, epochs: int, batches: int, verbose: bool):\n",
    "        if verbose:\n",
    "            print(f\"Log. reg. BGD step size:\", step)\n",
    "        \n",
    "        # Compute mini-batches\n",
    "        X_mini_batches = batch(X_tilde, batch_count=batches)\n",
    "        y_mini_batches = batch(np.expand_dims(y.flatten(), axis=1), batch_count=batches)\n",
    "        \n",
    "        # Mini-batched gradient descent\n",
    "        w = initial_w\n",
    "        for _ in range(epochs):\n",
    "            for i in range(batches):\n",
    "                if verbose:\n",
    "                    print(f\"Log. reg. MBGD epoch {_ + 1}/{epochs} -- batch {i + 1}/{batches}\")\n",
    "                    print(f\"\\t weigts avg {np.mean(w)} | loss: {CustomLogReg.loss(X_mini_batches[i], y_mini_batches[i], w)}\")\n",
    "                \n",
    "                dot = X_mini_batches[i] @ w\n",
    "                sig = sigmoid(dot)  # m x 1\n",
    "                w -= X_mini_batches[i].T @ (sig - y_mini_batches[i]) * step\n",
    "        \n",
    "        return w\n",
    "    \n",
    "    def log_loss(X, y, w):\n",
    "        a = sigmoid(X @ w)\n",
    "        return np.sum(-y * np.log(a) - (1 - y) * np.log(1 - a))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" PyMICE \"\"\"\n",
    "\n",
    "def pymice(incomplete_data, round = False):\n",
    "    imp_mean = IterativeImputer()  # set sample_posterior=True, random_state=random.randint(0, 1000) as attributes for MI\n",
    "    imp_mean.fit(incomplete_data)\n",
    "    imputed_data = imp_mean.transform(incomplete_data)\n",
    "\n",
    "    if round:\n",
    "        imputed_data[np.isnan(incomplete_data)] = np.round(imputed_data[np.isnan(incomplete_data)])\n",
    "    \n",
    "    return imputed_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Evaluators \"\"\"\n",
    "\n",
    "def sign(x):\n",
    "    return -1 if x < 0 else 1 if x > 0 else 0\n",
    "\n",
    "def get_stat_sig(X, y):\n",
    "    X2 = sm.add_constant(X)\n",
    "    est = sm.OLS(np.expand_dims(y, axis=1), X2)\n",
    "    est2 = est.fit()\n",
    "    return est2.params, est2.pvalues\n",
    "\n",
    "def discrepancies(data_instance, mdni_stats, sequre_stats, verbose=False):\n",
    "    mdni_coeff, mdni_pval = mdni_stats\n",
    "    sequre_coeff, sequre_pval = sequre_stats\n",
    "\n",
    "    discr = 0\n",
    "    for i in range(len(mdni_coeff)):\n",
    "        if mdni_pval[i] <= 0.05 and (sequre_pval[i] > 0.05 or sign(mdni_coeff[i]) != sign(sequre_coeff[i])):\n",
    "            if verbose:\n",
    "                print(f\"Discrepancy found at var {i} and instance {data_instance}! SIMICE p_val: {mdni_pval[i]}, Sequre p_val: {sequre_pval[i]}, SIMICE coeff: {mdni_coeff[i]}, Sequre coeff: {sequre_coeff[i]}\")\n",
    "            discr += 1\n",
    "        elif verbose:\n",
    "            print(f\"No discrepancy found at var {i} and instance {data_instance}:\\n\\t\\tSIMICE p_val: {mdni_pval[i]}\\n\\t\\tSequre p_val: {sequre_pval[i]}\\n\\t\\tSIMICE coeff: {mdni_coeff[i]}\\n\\t\\tSequre coeff: {sequre_coeff[i]}\")\n",
    "    \n",
    "    return discr\n",
    "\n",
    "def count_discrepancies(imputed_datasets, mice_imputed_data, labels):\n",
    "    mice_stats = get_stat_sig(mice_imputed_data, labels)\n",
    "\n",
    "    discr = 0\n",
    "    for i, imputed_dataset in enumerate(imputed_datasets):\n",
    "        sequre_stats = get_stat_sig(imputed_dataset, labels)\n",
    "        discr += discrepancies(i, mice_stats, sequre_stats)\n",
    "    \n",
    "    return discr / len(imputed_datasets)\n",
    "\n",
    "def evaluate_scenario(scenario, impute_model, mi_factor, miss_col, miss_val, dump_data):\n",
    "    (msg, tampered_data, labels, non_tampered_data, miss_rows) = scenario\n",
    "\n",
    "    mice_imputed_data = pymice(tampered_data, isinstance(impute_model, LogisticRegression))\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/{msg}/{len(tampered_data)}/py_mice_data.txt\", mice_imputed_data)\n",
    "    \n",
    "    fit_model = LinearRegression()\n",
    "    mi, imputed_datasets = MI(mi_factor, impute_model, fit_model).fit(\n",
    "        np.nan_to_num(tampered_data, nan=miss_val), labels, miss_rows, miss_col, noise=True if msg.endswith(\"_1\") else False)\n",
    "    \n",
    "    predicted_data = mi.model.predict(non_tampered_data).flatten()\n",
    "    labels = labels.flatten()\n",
    "    assert predicted_data.shape == labels.shape\n",
    "\n",
    "    abs_diff = np.abs(predicted_data - labels)\n",
    "    y_hat_y_mean = np.mean(abs_diff)\n",
    "    y_hat_y_std = np.std(abs_diff)\n",
    "    discrepancies_count = count_discrepancies(imputed_datasets, mice_imputed_data, labels)\n",
    "\n",
    "    imputed_cols = [imputed_dataset[miss_rows, miss_col] for imputed_dataset in imputed_datasets]\n",
    "    imputed_cols_prob = copy(imputed_cols)\n",
    "    if \"_2\" in msg:\n",
    "        imputed_cols = [np.round(imputed_col) for imputed_col in imputed_cols]\n",
    "    expected_imputation = non_tampered_data[miss_rows, miss_col]\n",
    "    imputation_abs_diff = sum([np.abs(imputed_col - expected_imputation) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "    imputation_diff_mean = np.mean(imputation_abs_diff)\n",
    "    imputation_diff_std = np.std(imputation_abs_diff)\n",
    "    imputation_accuracy = sum([(np.sum(imputed_col == expected_imputation) / len(expected_imputation)) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "    imputation_auc = sum([roc_auc_score(expected_imputation, imputed_col_prob) for imputed_col_prob in imputed_cols_prob]) / len(imputed_cols_prob)\n",
    "\n",
    "    return mi.model.coef_, y_hat_y_mean, y_hat_y_std, discrepancies_count, imputation_diff_mean, imputation_diff_std, imputation_accuracy, imputation_auc\n",
    "\n",
    "def benchmark_scenario(runs_count, scenario, impute_model, mi_factor, miss_col, miss_val, dump_data):\n",
    "    features_count = scenario[1].shape[1]\n",
    "    \n",
    "    coefs = np.zeros((runs_count, features_count))\n",
    "    y_hat_y_means = np.zeros((runs_count,))\n",
    "    y_hat_y_sds = np.zeros((runs_count,))\n",
    "    discrs = np.zeros((runs_count,))\n",
    "    imp_means = np.zeros((runs_count,))\n",
    "    imp_sds = np.zeros((runs_count,))\n",
    "    imp_accs = np.zeros((runs_count,))\n",
    "    imp_aucs = np.zeros((runs_count,))\n",
    "    runtimes = np.zeros((runs_count,))\n",
    "    \n",
    "    for i in tqdm(range(runs_count), f\"Evaluating {scenario[0]}\"):\n",
    "        s = time.time()\n",
    "        coef, y_hat_y_mean, y_hat_y_sd, discr, imp_mean, imp_sd, imp_acc, imp_auc = evaluate_scenario(\n",
    "            scenario, impute_model, mi_factor, miss_col, miss_val, dump_data)\n",
    "        e = time.time()\n",
    "        \n",
    "        coefs[i] = coef\n",
    "        y_hat_y_means[i] = y_hat_y_mean\n",
    "        y_hat_y_sds[i] = y_hat_y_sd\n",
    "        discrs[i] = discr\n",
    "        imp_means[i] = imp_mean\n",
    "        imp_sds[i] = imp_sd\n",
    "        imp_accs[i] = imp_acc\n",
    "        imp_aucs[i] = imp_auc\n",
    "        runtimes[i] = e - s\n",
    "    \n",
    "    return scenario[0], coefs, y_hat_y_means, y_hat_y_sds, discrs, imp_means, imp_sds, imp_accs, imp_aucs, runtimes\n",
    "\n",
    "def evaluate_real_setup(real_data, mi_factor, miss_val, dump_data):\n",
    "    incomplete_data, complete_data, labels, miss_rows, col_types, non_miss_rows = real_data\n",
    "\n",
    "    impute_models = [LogisticRegression() if t == \"log\" else LinearRegression() for t in col_types]\n",
    "    fit_model = LinearRegression()\n",
    "    \n",
    "    mice_imputed_data = pymice(incomplete_data)\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/real/py_mice_data.txt\", mice_imputed_data, fmt='%10.2f')\n",
    "\n",
    "    incompletes_w_miss_val = np.nan_to_num(incomplete_data, nan=miss_val)\n",
    "    mi, imputed_datasets = BatchMI(mi_factor, impute_models, fit_model).fit(\n",
    "        incompletes_w_miss_val, complete_data, labels, miss_rows)\n",
    "    \n",
    "    predicted_data = mi.model.predict(complete_data).flatten()\n",
    "    complete_labels = labels.flatten()[non_miss_rows]\n",
    "    assert predicted_data.shape == complete_labels.shape\n",
    "\n",
    "    abs_diff = np.abs(predicted_data - complete_labels)\n",
    "    y_hat_y_mean = np.mean(abs_diff)\n",
    "    y_hat_y_sd = np.std(abs_diff)\n",
    "    discrepancies_count = count_discrepancies(imputed_datasets, mice_imputed_data, labels.flatten())\n",
    "\n",
    "    return mi.model.coef_, y_hat_y_mean, y_hat_y_sd, discrepancies_count\n",
    "\n",
    "def benchmark_real_setup(runs_count, real_data, mi_factor, miss_val, dump_data):\n",
    "    y_hat_y_means = np.zeros((runs_count,))\n",
    "    y_hat_y_sds = np.zeros((runs_count,))\n",
    "    discrs = np.zeros((runs_count,))\n",
    "    runtimes = np.zeros((runs_count,))\n",
    "    \n",
    "    for i in tqdm(range(runs_count), \"Evaluating real setup\"):\n",
    "        s = time.time()\n",
    "        _, y_hat_y_mean, y_hat_y_sd, discr = evaluate_real_setup(real_data, mi_factor, miss_val, dump_data)\n",
    "        e = time.time()\n",
    "\n",
    "        y_hat_y_means[i] = y_hat_y_mean\n",
    "        y_hat_y_sds[i] = y_hat_y_sd\n",
    "        discrs[i] = discr\n",
    "        runtimes[i] = e - s\n",
    "    \n",
    "    return None, y_hat_y_means, y_hat_y_sds, discrs, None, None, None, None, runtimes\n",
    "    \n",
    "def evaluate_pymice(data, labels, non_tampered_data, mi_factor, non_miss_rows = None, binary = False):\n",
    "    fit_model = LinearRegression()\n",
    "    imputed_datasets = [pymice(data) for _ in range(mi_factor)]\n",
    "\n",
    "    mi = MI(mi_factor, None, fit_model)\n",
    "    mi.fit_analysis_model(imputed_datasets, labels)\n",
    "\n",
    "    if binary:\n",
    "        for imputed_dataset in imputed_datasets:\n",
    "            imputed_dataset[np.isnan(data)] = np.round(imputed_dataset[np.isnan(data)])\n",
    "    \n",
    "    predicted_data = mi.model.predict(non_tampered_data).flatten()\n",
    "    labels = labels.flatten()\n",
    "    if non_miss_rows is not None:\n",
    "        labels = labels[non_miss_rows]\n",
    "    assert predicted_data.shape == labels.shape\n",
    "\n",
    "    abs_diff = np.abs(predicted_data - labels)\n",
    "    y_hat_y_mean = np.mean(abs_diff)\n",
    "    y_hat_y_sds = np.std(abs_diff)\n",
    "\n",
    "    if non_miss_rows is None:\n",
    "        imputed_cols = [imputed_dataset[np.where(np.isnan(data))] for imputed_dataset in imputed_datasets]\n",
    "        imputed_cols_prob = copy(imputed_cols)\n",
    "        if binary:\n",
    "            imputed_cols = [np.round(imputed_col) for imputed_col in imputed_cols]\n",
    "        expected_imputation = non_tampered_data[np.where(np.isnan(data))]\n",
    "        imputation_abs_diff = sum([np.abs(imputed_col - expected_imputation) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "        imputation_diff_mean = np.mean(imputation_abs_diff)\n",
    "        imputation_diff_std = np.std(imputation_abs_diff)\n",
    "        imputation_accuracy = sum([(np.sum(imputed_col == expected_imputation) / len(expected_imputation)) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "        imputation_auc = sum([roc_auc_score(expected_imputation, imputed_col_prob) for imputed_col_prob in imputed_cols_prob]) / len(imputed_cols_prob)\n",
    "\n",
    "        return mi.model.coef_, y_hat_y_mean, y_hat_y_sds, imputation_diff_mean, imputation_diff_std, imputation_accuracy, imputation_auc\n",
    "    \n",
    "    return mi.model.coef_, y_hat_y_mean, y_hat_y_sds, None, None, None\n",
    "\n",
    "def benchmark_pymice(runs_count, msg, data, labels, non_tampered_data, mi_factor, non_miss_rows = None, binary = False):\n",
    "    coefs = np.zeros((runs_count, data.shape[1]))\n",
    "    y_hat_y_means = np.zeros((runs_count,))\n",
    "    y_hat_y_sds = np.zeros((runs_count,))\n",
    "    discrs = np.zeros((runs_count,)) - 1\n",
    "    imp_means = np.zeros((runs_count,))\n",
    "    imp_sds = np.zeros((runs_count,))\n",
    "    imp_accs = np.zeros((runs_count,))\n",
    "    imp_aucs = np.zeros((runs_count,))\n",
    "    runtimes = np.zeros((runs_count,))\n",
    "    \n",
    "    for i in tqdm(range(runs_count), f\"Evaluating pymice for {msg}\"):\n",
    "        s = time.time()\n",
    "        coef, y_hat_y_mean, y_hat_y_sd, imp_mean, imp_sd, imp_acc, imp_auc = evaluate_pymice(data, labels, non_tampered_data, mi_factor, non_miss_rows, binary)\n",
    "        e = time.time()\n",
    "        \n",
    "        coefs[i] = coef\n",
    "        y_hat_y_means[i] = y_hat_y_mean\n",
    "        y_hat_y_sds[i] = y_hat_y_sd\n",
    "        imp_means[i] = imp_mean\n",
    "        imp_sds[i] = imp_sd\n",
    "        imp_accs[i] = imp_acc\n",
    "        imp_aucs[i] = imp_auc\n",
    "        runtimes[i] = e - s\n",
    "    \n",
    "    return msg, coefs, y_hat_y_means, y_hat_y_sds, discrs, imp_means, imp_sds, imp_accs, imp_aucs, runtimes\n",
    "\n",
    "def evaluate_simice(scenario, tampered_data, labels, non_tampered_data, mi_factor, dump_data):\n",
    "    rows = tampered_data.shape[0]\n",
    "    \n",
    "    fit_model = LinearRegression()\n",
    "    imputed_datasets = [np.loadtxt(f\"../../data/mi/siMICE/{scenario}/{rows}/imputed_data_{i + 1}.txt\") for i in range(mi_factor)]\n",
    "    mice_imputed_data = pymice(tampered_data, \"_2\" in scenario)\n",
    "    if dump_data:\n",
    "        np.savetxt(f\"../../data/mi/siMICE/{scenario}/{rows}/py_mice_data.txt\", mice_imputed_data, fmt='%10.2f')\n",
    "    \n",
    "    mi = MI(mi_factor, None, fit_model)\n",
    "    mi.fit_analysis_model(imputed_datasets, labels)\n",
    "    \n",
    "    predicted_data = mi.model.predict(non_tampered_data).flatten()\n",
    "    labels = labels.flatten()\n",
    "    assert predicted_data.shape == labels.shape\n",
    "    \n",
    "    abs_diff = np.abs(predicted_data - labels)\n",
    "    y_hat_y_mean = np.mean(abs_diff)\n",
    "    y_hat_y_sds = np.std(abs_diff)\n",
    "    discrepancies_count = count_discrepancies(imputed_datasets, mice_imputed_data, labels)\n",
    "\n",
    "    imputed_cols = [imputed_dataset[np.where(np.isnan(tampered_data))] for imputed_dataset in imputed_datasets]\n",
    "    imputed_cols_prob = copy(imputed_cols)\n",
    "    if \"_2\" in scenario:\n",
    "        imputed_cols = [np.round(imputed_col) for imputed_col in imputed_cols]\n",
    "    expected_imputation = non_tampered_data[np.where(np.isnan(tampered_data))]\n",
    "    imputation_abs_diff = sum([np.abs(imputed_col - expected_imputation) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "    imputation_diff_mean = np.mean(imputation_abs_diff)\n",
    "    imputation_diff_std = np.std(imputation_abs_diff)\n",
    "    imputation_accuracy = sum([(np.sum(imputed_col == expected_imputation) / len(expected_imputation)) for imputed_col in imputed_cols]) / len(imputed_cols)\n",
    "    imputation_auc = sum([roc_auc_score(expected_imputation, imputed_col_prob) for imputed_col_prob in imputed_cols_prob]) / len(imputed_cols_prob)\n",
    "\n",
    "    return mi.model.coef_, y_hat_y_mean, y_hat_y_sds, discrepancies_count, imputation_diff_mean, imputation_diff_std, imputation_accuracy, imputation_auc\n",
    "\n",
    "def benchmark_simice(runs_count, scenario, tampered_data, labels, non_tampered_data, mi_factor, dump_data):\n",
    "    cols = tampered_data.shape[1]\n",
    "    coefs = np.zeros((runs_count, cols))\n",
    "    y_hat_y_means = np.zeros((runs_count,))\n",
    "    y_hat_y_sds = np.zeros((runs_count,))\n",
    "    discrs = np.zeros((runs_count,))\n",
    "    imp_means = np.zeros((runs_count,))\n",
    "    imp_sds = np.zeros((runs_count,))\n",
    "    imp_accs = np.zeros((runs_count,))\n",
    "    imp_aucs = np.zeros((runs_count,))\n",
    "    runtimes = np.zeros((runs_count,))\n",
    "    \n",
    "    for i in tqdm(range(runs_count), \"Evaluating siMICE\"):\n",
    "        s = time.time()\n",
    "        coef, y_hat_y_mean, y_hat_y_sd, discr, imp_mean, imp_sd, imp_acc, imp_auc = evaluate_simice(scenario, tampered_data, labels, non_tampered_data, mi_factor, dump_data)\n",
    "        e = time.time()\n",
    "\n",
    "        coefs[i] = coef\n",
    "        y_hat_y_means[i] = y_hat_y_mean\n",
    "        y_hat_y_sds[i] = y_hat_y_sd\n",
    "        discrs[i] = discr\n",
    "        imp_means[i] = imp_mean\n",
    "        imp_sds[i] = imp_sd\n",
    "        imp_accs[i] = imp_acc\n",
    "        imp_aucs[i] = imp_auc\n",
    "        runtimes[i] = e - s\n",
    "    \n",
    "    return coefs, y_hat_y_means, y_hat_y_sds, discrs, imp_means, imp_sds, imp_accs, imp_aucs, runtimes\n",
    "\n",
    "def print_stats(msg, coefs, biases, sds, discrepancies_count, imp_biases, imp_sds, imp_accs, imp_aucs, runtimes):\n",
    "    csv_out_path = \"temp_mi_stats_python.csv\"\n",
    "    \n",
    "    with open(csv_out_path, \"a+\") as f:\n",
    "        if msg.lower().startswith(\"midn\") or msg.lower().startswith(\"paper\"):\n",
    "            coefs_truth = np.ones_like(coefs[0])\n",
    "            coefs_bias = np.linalg.norm(np.mean(coefs, axis=0) - coefs_truth)\n",
    "            coefs_sd = np.sqrt(np.mean(np.sum((coefs - np.mean(coefs, axis=0)) ** 2, axis=1)))\n",
    "            coefs_rmse = np.sqrt(np.mean(np.sum((coefs - coefs_truth) ** 2, axis=1)))\n",
    "\n",
    "            print(f\"{msg} | Python MI coefs bias: {coefs_bias}\")\n",
    "            print(f\"{msg} | Python MI coefs SD: {coefs_sd}\")\n",
    "            print(f\"{msg} | Python MI coefs rMSE: {coefs_rmse}\")\n",
    "\n",
    "            f.write(f\"{coefs_bias},{coefs_sd},{coefs_rmse},\")\n",
    "        else:\n",
    "            f.write(f\",,,\")\n",
    "    \n",
    "        print(f\"{msg} | Python MI runtime (s) | mean: {np.mean(runtimes)}, min: {np.min(runtimes)}, max: {np.max(runtimes)}\")\n",
    "        print(f\"{msg} | Python MI y_hat - y mean | mean: {np.mean(biases)}, min: {np.min(biases)}, max: {np.max(biases)}\")\n",
    "        print(f\"{msg} | Python MI y_hat - y SD | mean: {np.mean(sds)}, min: {np.min(sds)}, max: {np.max(sds)}\")\n",
    "        print(f\"{msg} | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: {np.mean(discrepancies_count)}, min: {np.min(discrepancies_count)}, max: {np.max(discrepancies_count)}\")\n",
    "\n",
    "        f.write(f\"{np.mean(biases)},{np.mean(sds)},{np.mean(discrepancies_count)},\")\n",
    "    \n",
    "        if imp_biases is not None:\n",
    "            print(f\"{msg} | Python MI imputation diff mean | mean: {np.mean(imp_biases)}, min: {np.min(imp_biases)}, max: {np.max(imp_biases)}\")\n",
    "            f.write(f\"{np.mean(imp_biases)},\")\n",
    "        else:\n",
    "            f.write(\",\")\n",
    "        if imp_sds is not None:\n",
    "            print(f\"{msg} | Python MI imputation diff SD | mean: {np.mean(imp_sds)}, min: {np.min(imp_sds)}, max: {np.max(imp_sds)}\")\n",
    "            f.write(f\"{np.mean(imp_sds)},\")\n",
    "        else:\n",
    "            f.write(\",\")\n",
    "        if imp_accs is not None:\n",
    "            print(f\"{msg} | Python MI imputation accuracy | mean: {np.mean(imp_accs)}, min: {np.min(imp_accs)}, max: {np.max(imp_accs)}\")\n",
    "            f.write(f\"{np.mean(imp_accs)},\")\n",
    "        else:\n",
    "            f.write(\",\")\n",
    "        if imp_aucs is not None:\n",
    "            print(f\"{msg} | Python MI imputation AUC | mean: {np.mean(imp_aucs)}, min: {np.min(imp_aucs)}, max: {np.max(imp_aucs)}\")\n",
    "            f.write(f\"{np.mean(imp_aucs)},\")\n",
    "        else:\n",
    "            f.write(\",\")\n",
    "    \n",
    "        f.write(f\"{np.mean(runtimes)}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data generated by siMICE scripts (https://github.com/Luyaochen1/midn_gui)\n",
    "midn_sc1 = load_or_generate_midn_data(midn_scenario_rows, midn_scenario_cols, \"scenario_1\", midn_miss_rows_count, scenarios_miss_col, labels_noise_scale)\n",
    "midn_sc2 = load_or_generate_midn_data(midn_scenario_rows, midn_scenario_cols, \"scenario_2\", midn_miss_rows_count, scenarios_miss_col, labels_noise_scale)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated data missing rate: 2438/5000\n",
      "Generated data missing rate: 2490/5000\n"
     ]
    }
   ],
   "source": [
    "# Scenarios per data generation instructions in https://www.nature.com/articles/s41467-020-19270-2 \n",
    "paper_sc1 = load_or_generate_paper_data(paper_scenario_rows, \"scenario_1\", scenarios_miss_col, labels_noise_scale)\n",
    "paper_sc2 = load_or_generate_paper_data(paper_scenario_rows, \"scenario_2\", scenarios_miss_col, labels_noise_scale)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_real_setup:\n",
    "    real = load_real_data(real_setup_outlier_threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imputation via regression analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating midn_scenario_1: 100%|██████████| 100/100 [00:14<00:00,  6.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "midn_scenario_1 | Python MI coefs bias: 0.015706849934592572\n",
      "midn_scenario_1 | Python MI coefs SD: 0.00014585079700028537\n",
      "midn_scenario_1 | Python MI coefs rMSE: 0.015707527091264342\n",
      "midn_scenario_1 | Python MI runtime (s) | mean: 0.14504011869430541, min: 0.1364457607269287, max: 0.16270232200622559\n",
      "midn_scenario_1 | Python MI y_hat - y mean | mean: 0.025579808819667118, min: 0.02553802624295158, max: 0.025623922390603456\n",
      "midn_scenario_1 | Python MI y_hat - y SD | mean: 0.02573941093419772, min: 0.02571982155562864, max: 0.025759605006982423\n",
      "midn_scenario_1 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: 0.0, min: 0.0, max: 0.0\n",
      "midn_scenario_1 | Python MI imputation diff mean | mean: 0.5008954533649917, min: 0.5005325699977577, max: 0.5012015348204785\n",
      "midn_scenario_1 | Python MI imputation diff SD | mean: 0.03229444188389131, min: 0.031986050557552395, max: 0.03251902741722119\n",
      "midn_scenario_1 | Python MI imputation accuracy | mean: 0.0, min: 0.0, max: 0.0\n",
      "midn_scenario_1 | Python MI imputation AUC | mean: 0.4863912111277833, min: 0.4811784868547372, max: 0.4930810476386884\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating midn_scenario_2: 100%|██████████| 100/100 [00:21<00:00,  4.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "midn_scenario_2 | Python MI coefs bias: 0.012120464650858491\n",
      "midn_scenario_2 | Python MI coefs SD: 3.6905523389740094e-15\n",
      "midn_scenario_2 | Python MI coefs rMSE: 0.012120464650858503\n",
      "midn_scenario_2 | Python MI runtime (s) | mean: 0.21338937044143677, min: 0.19252991676330566, max: 0.2425227165222168\n",
      "midn_scenario_2 | Python MI y_hat - y mean | mean: 0.024399518728786165, min: 0.024399518728786172, max: 0.024399518728786172\n",
      "midn_scenario_2 | Python MI y_hat - y SD | mean: 0.025504167190784008, min: 0.025504167190784, max: 0.025504167190784\n",
      "midn_scenario_2 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: 0.0, min: 0.0, max: 0.0\n",
      "midn_scenario_2 | Python MI imputation diff mean | mean: 0.4879999999999998, min: 0.488, max: 0.488\n",
      "midn_scenario_2 | Python MI imputation diff SD | mean: 0.4998559792580257, min: 0.4998559792580259, max: 0.4998559792580259\n",
      "midn_scenario_2 | Python MI imputation accuracy | mean: 0.5120000000000001, min: 0.512, max: 0.512\n",
      "midn_scenario_2 | Python MI imputation AUC | mean: 0.5095126479626676, min: 0.5095126479626678, max: 0.5095126479626678\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating paper_scenario_1: 100%|██████████| 100/100 [00:09<00:00, 11.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "paper_scenario_1 | Python MI coefs bias: 0.00235847166366012\n",
      "paper_scenario_1 | Python MI coefs SD: 0.00031236771146818116\n",
      "paper_scenario_1 | Python MI coefs rMSE: 0.002379067501239456\n",
      "paper_scenario_1 | Python MI runtime (s) | mean: 0.09063749074935913, min: 0.08864665031433105, max: 0.12207269668579102\n",
      "paper_scenario_1 | Python MI y_hat - y mean | mean: 0.021643935191931342, min: 0.02129994506316016, max: 0.022026030425999224\n",
      "paper_scenario_1 | Python MI y_hat - y SD | mean: 0.013843577574287457, min: 0.013735568006013112, max: 0.013964296570760336\n",
      "paper_scenario_1 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: 0.0, min: 0.0, max: 0.0\n",
      "paper_scenario_1 | Python MI imputation diff mean | mean: 0.4947428759453218, min: 0.49448270298145075, max: 0.4949339924455822\n",
      "paper_scenario_1 | Python MI imputation diff SD | mean: 0.07420115775669307, min: 0.07395804863982766, max: 0.07442153758899597\n",
      "paper_scenario_1 | Python MI imputation accuracy | mean: 0.0, min: 0.0, max: 0.0\n",
      "paper_scenario_1 | Python MI imputation AUC | mean: 0.5218651130554144, min: 0.5176491012851137, max: 0.5274055349596308\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating paper_scenario_2: 100%|██████████| 100/100 [00:22<00:00,  4.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "paper_scenario_2 | Python MI coefs bias: 0.08654131160342107\n",
      "paper_scenario_2 | Python MI coefs SD: 1.159106867033638e-15\n",
      "paper_scenario_2 | Python MI coefs rMSE: 0.08654131160342224\n",
      "paper_scenario_2 | Python MI runtime (s) | mean: 0.22842355966567993, min: 0.22569060325622559, max: 0.24114489555358887\n",
      "paper_scenario_2 | Python MI y_hat - y mean | mean: 0.02221221024418168, min: 0.022212210244181678, max: 0.022212210244181678\n",
      "paper_scenario_2 | Python MI y_hat - y SD | mean: 0.015207886977498253, min: 0.015207886977498248, max: 0.015207886977498248\n",
      "paper_scenario_2 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: 0.0, min: 0.0, max: 0.0\n",
      "paper_scenario_2 | Python MI imputation diff mean | mean: 0.4875502008032129, min: 0.4875502008032129, max: 0.4875502008032129\n",
      "paper_scenario_2 | Python MI imputation diff SD | mean: 0.4998449784682844, min: 0.49984497846828435, max: 0.49984497846828435\n",
      "paper_scenario_2 | Python MI imputation accuracy | mean: 0.5124497991967871, min: 0.5124497991967871, max: 0.5124497991967871\n",
      "paper_scenario_2 | Python MI imputation AUC | mean: 0.5103018696404709, min: 0.5103018696404706, max: 0.5103018696404706\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "if run_scenarios:\n",
    "    print_stats(*benchmark_scenario(scenarios_runs_count, midn_sc1, LinearRegression(), mi_factor, scenarios_miss_col, miss_val, dump_data))\n",
    "    print_stats(*benchmark_scenario(scenarios_runs_count, midn_sc2, LogisticRegression(), mi_factor, scenarios_miss_col, miss_val, dump_data))\n",
    "    print_stats(*benchmark_scenario(scenarios_runs_count, paper_sc1, LinearRegression(), mi_factor, scenarios_miss_col, miss_val, dump_data))\n",
    "    print_stats(*benchmark_scenario(scenarios_runs_count, paper_sc2, LogisticRegression(), mi_factor, scenarios_miss_col, miss_val, dump_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyMICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating pymice for midn_scenario_1: 100%|██████████| 100/100 [00:20<00:00,  4.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "midn_scenario_1 | Python MI coefs bias: 0.011217878813363068\n",
      "midn_scenario_1 | Python MI coefs SD: 4.429775955663296e-15\n",
      "midn_scenario_1 | Python MI coefs rMSE: 0.01121787881336602\n",
      "midn_scenario_1 | Python MI runtime (s) | mean: 0.2017705774307251, min: 0.20020699501037598, max: 0.2157878875732422\n",
      "midn_scenario_1 | Python MI y_hat - y mean | mean: 0.02393902088218626, min: 0.023939020882186263, max: 0.023939020882186263\n",
      "midn_scenario_1 | Python MI y_hat - y SD | mean: 0.02511140545677966, min: 0.02511140545677967, max: 0.02511140545677967\n",
      "midn_scenario_1 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: -1.0, min: -1.0, max: -1.0\n",
      "midn_scenario_1 | Python MI imputation diff mean | mean: 0.5003563775707297, min: 0.5003563775707297, max: 0.5003563775707297\n",
      "midn_scenario_1 | Python MI imputation diff SD | mean: 0.015037819745451906, min: 0.015037819745451908, max: 0.015037819745451908\n",
      "midn_scenario_1 | Python MI imputation accuracy | mean: 0.0, min: 0.0, max: 0.0\n",
      "midn_scenario_1 | Python MI imputation AUC | mean: 0.4864458366815768, min: 0.4864458366815767, max: 0.4864458366815767\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating pymice for midn_scenario_2: 100%|██████████| 100/100 [00:20<00:00,  4.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "midn_scenario_2 | Python MI coefs bias: 0.012919182964483253\n",
      "midn_scenario_2 | Python MI coefs SD: 5.453670444495245e-15\n",
      "midn_scenario_2 | Python MI coefs rMSE: 0.012919182964485451\n",
      "midn_scenario_2 | Python MI runtime (s) | mean: 0.20405047178268432, min: 0.20147061347961426, max: 0.22124171257019043\n",
      "midn_scenario_2 | Python MI y_hat - y mean | mean: 0.02496504948282109, min: 0.02496504948282108, max: 0.02496504948282108\n",
      "midn_scenario_2 | Python MI y_hat - y SD | mean: 0.025692030736707127, min: 0.025692030736707124, max: 0.025692030736707124\n",
      "midn_scenario_2 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: -1.0, min: -1.0, max: -1.0\n",
      "midn_scenario_2 | Python MI imputation diff mean | mean: 0.5126666666666668, min: 0.5126666666666667, max: 0.5126666666666667\n",
      "midn_scenario_2 | Python MI imputation diff SD | mean: 0.499839529804872, min: 0.4998395298048721, max: 0.4998395298048721\n",
      "midn_scenario_2 | Python MI imputation accuracy | mean: 0.4873333333333332, min: 0.4873333333333333, max: 0.4873333333333333\n",
      "midn_scenario_2 | Python MI imputation AUC | mean: 0.49620489984065563, min: 0.4962048998406556, max: 0.4962048998406556\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating pymice for paper_scenario_1: 100%|██████████| 100/100 [00:03<00:00, 27.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "paper_scenario_1 | Python MI coefs bias: 0.015076296995715435\n",
      "paper_scenario_1 | Python MI coefs SD: 1.9389211565826713e-15\n",
      "paper_scenario_1 | Python MI coefs rMSE: 0.015076296995714547\n",
      "paper_scenario_1 | Python MI runtime (s) | mean: 0.036411161422729495, min: 0.03572368621826172, max: 0.041876792907714844\n",
      "paper_scenario_1 | Python MI y_hat - y mean | mean: 0.021649977746235564, min: 0.02164997774623557, max: 0.02164997774623557\n",
      "paper_scenario_1 | Python MI y_hat - y SD | mean: 0.014056423955094315, min: 0.014056423955094316, max: 0.014056423955094316\n",
      "paper_scenario_1 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: -1.0, min: -1.0, max: -1.0\n",
      "paper_scenario_1 | Python MI imputation diff mean | mean: 0.4949256194260755, min: 0.49492561942607555, max: 0.49492561942607555\n",
      "paper_scenario_1 | Python MI imputation diff SD | mean: 0.0729580541588863, min: 0.0729580541588863, max: 0.0729580541588863\n",
      "paper_scenario_1 | Python MI imputation accuracy | mean: 0.0, min: 0.0, max: 0.0\n",
      "paper_scenario_1 | Python MI imputation AUC | mean: 0.524395137320176, min: 0.524395137320176, max: 0.524395137320176\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating pymice for paper_scenario_2: 100%|██████████| 100/100 [00:03<00:00, 27.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "paper_scenario_2 | Python MI coefs bias: 0.006500942206016211\n",
      "paper_scenario_2 | Python MI coefs SD: 1.2947314098277875e-15\n",
      "paper_scenario_2 | Python MI coefs rMSE: 0.006500942206016647\n",
      "paper_scenario_2 | Python MI runtime (s) | mean: 0.036467111110687254, min: 0.0359950065612793, max: 0.04524683952331543\n",
      "paper_scenario_2 | Python MI y_hat - y mean | mean: 0.02493872832530602, min: 0.024938728325306022, max: 0.024938728325306022\n",
      "paper_scenario_2 | Python MI y_hat - y SD | mean: 0.014006148255582485, min: 0.014006148255582484, max: 0.014006148255582484\n",
      "paper_scenario_2 | Python MI # of discrepancies w.r.t. PyMICE per imputed dataset | mean: -1.0, min: -1.0, max: -1.0\n",
      "paper_scenario_2 | Python MI imputation diff mean | mean: 0.4767068273092371, min: 0.47670682730923697, max: 0.47670682730923697\n",
      "paper_scenario_2 | Python MI imputation diff SD | mean: 0.49945713340185477, min: 0.4994571334018549, max: 0.4994571334018549\n",
      "paper_scenario_2 | Python MI imputation accuracy | mean: 0.5232931726907629, min: 0.523293172690763, max: 0.523293172690763\n",
      "paper_scenario_2 | Python MI imputation AUC | mean: 0.5, min: 0.5, max: 0.5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "if run_pymice:\n",
    "    print_stats(*benchmark_pymice(pymice_runs_count, midn_sc1[0], midn_sc1[1], midn_sc1[2], midn_sc1[3], mi_factor))\n",
    "    print_stats(*benchmark_pymice(pymice_runs_count, midn_sc2[0], midn_sc2[1], midn_sc2[2], midn_sc2[3], mi_factor, binary=True))\n",
    "    print_stats(*benchmark_pymice(pymice_runs_count, paper_sc1[0], paper_sc1[1], paper_sc1[2], paper_sc1[3], mi_factor))\n",
    "    print_stats(*benchmark_pymice(pymice_runs_count, paper_sc2[0], paper_sc2[1], paper_sc2[2], paper_sc2[3], mi_factor, binary=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### siMICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_simice:\n",
    "    print_stats(\"siMICE scenario 1\", *benchmark_simice(simice_runs_count, \"scenario_1\", midn_sc1[1], midn_sc1[2], midn_sc1[3], mi_factor, dump_data))\n",
    "    print_stats(\"siMICE scenario 2\", *benchmark_simice(simice_runs_count, \"scenario_2\", midn_sc2[0], midn_sc2[2], midn_sc2[3], mi_factor, dump_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "if run_real_setup:\n",
    "    print_stats(\"Real data\", *benchmark_real_setup(real_setup_runs_count, real, mi_factor, miss_val, dump_data))\n",
    "    print_stats(\"PyMICE real setup\", *benchmark_pymice(real_setup_runs_count, real[0], real[1], real[3], real[2], mi_factor, non_miss_rows=real[-1]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}