notebooks/cohort/cohort_manager_scenarios.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cohort Manager - Scenarios and Examples\n",
    "\n",
    "The ``CohortManager`` class can be used in different ways to target mitigations to different cohorts. The main differences between these scenarios consist on whether the same or different type of data mitigation is applied to the cohort data, and whether a single or separate models will be trained for different cohorts. Depending on these choices, CohortManager will take care of slicing the data accordingly, applying the specified data mitigation strategy, merging the data back, and retraining the model(s). These different scenarios are summarized in the image below.\n",
    "\n",
    "![scenarios](./imgs/scenarios.jpg)\n",
    "\n",
    "In this notebook, we'll show a code snippet for each of the scenarios depicted in the image above. For an in-depth tutorial of how to use the ``CohortManager``, please check the [Managing Cohorts](./cohort_manager.ipynb) notebook.\n",
    "\n",
    "First of all, let's define some variables and create our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from lightgbm import LGBMClassifier\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "from raimitigations.utils import create_dummy_dataset, split_data\n",
    "import raimitigations.dataprocessing as dp\n",
    "from raimitigations.cohort import CohortDefinition, CohortManager\n",
    "\n",
    "SEED = 42\n",
    "\n",
    "np.random.seed(SEED)\n",
    "random.seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_df(with_null: bool = True):\n",
    "    np.random.seed(SEED)\n",
    "    random.seed(SEED)\n",
    "    def add_nan(vec, pct):\n",
    "        vec = list(vec)\n",
    "        nan_index = random.sample(range(len(vec)), int(pct * len(vec)))\n",
    "        for index in nan_index:\n",
    "            vec[index] = np.nan\n",
    "        return vec\n",
    "\n",
    "    df = create_dummy_dataset(\n",
    "        samples=1000,\n",
    "        n_features=5,\n",
    "        n_num_num=0,\n",
    "        n_cat_num=0,\n",
    "        n_cat_cat=0,\n",
    "        num_num_noise=[0.01, 0.05],\n",
    "        pct_change=[0.05, 0.1],\n",
    "    )\n",
    "    if with_null:\n",
    "        col_with_nan = [\"num_0\"]\n",
    "        for col in col_with_nan:\n",
    "            if col != \"label\":\n",
    "                df[col] = add_nan(df[col], 0.1)\n",
    "\n",
    "    return df\n",
    "\n",
    "# -----------------------------------\n",
    "def get_model():\n",
    "    model = LGBMClassifier(random_state=SEED)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>num_0</th>\n",
       "      <th>num_1</th>\n",
       "      <th>num_2</th>\n",
       "      <th>num_3</th>\n",
       "      <th>num_4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>131</th>\n",
       "      <td>2.945188</td>\n",
       "      <td>-2.967287</td>\n",
       "      <td>2.337389</td>\n",
       "      <td>1.799443</td>\n",
       "      <td>1.135685</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>333</th>\n",
       "      <td>2.974430</td>\n",
       "      <td>3.537706</td>\n",
       "      <td>1.987065</td>\n",
       "      <td>-1.861848</td>\n",
       "      <td>0.262459</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>211</th>\n",
       "      <td>1.353304</td>\n",
       "      <td>-3.212935</td>\n",
       "      <td>1.034578</td>\n",
       "      <td>1.797846</td>\n",
       "      <td>1.682174</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>850</th>\n",
       "      <td>1.205781</td>\n",
       "      <td>2.802149</td>\n",
       "      <td>4.308715</td>\n",
       "      <td>-1.498583</td>\n",
       "      <td>2.202188</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>514</th>\n",
       "      <td>1.659687</td>\n",
       "      <td>-2.176888</td>\n",
       "      <td>2.090389</td>\n",
       "      <td>1.608092</td>\n",
       "      <td>1.805161</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>95</th>\n",
       "      <td>NaN</td>\n",
       "      <td>3.770849</td>\n",
       "      <td>1.538188</td>\n",
       "      <td>-3.086773</td>\n",
       "      <td>-0.145386</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>317</th>\n",
       "      <td>2.553215</td>\n",
       "      <td>3.059551</td>\n",
       "      <td>2.469591</td>\n",
       "      <td>-2.561048</td>\n",
       "      <td>2.554759</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>481</th>\n",
       "      <td>1.915074</td>\n",
       "      <td>-2.785880</td>\n",
       "      <td>2.879047</td>\n",
       "      <td>-0.472673</td>\n",
       "      <td>5.318281</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>140</th>\n",
       "      <td>2.302969</td>\n",
       "      <td>3.789408</td>\n",
       "      <td>0.752376</td>\n",
       "      <td>-2.758305</td>\n",
       "      <td>1.449600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>291</th>\n",
       "      <td>-0.034102</td>\n",
       "      <td>-0.700874</td>\n",
       "      <td>0.013272</td>\n",
       "      <td>2.687944</td>\n",
       "      <td>4.080553</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>800 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "        num_0     num_1     num_2     num_3     num_4\n",
       "131  2.945188 -2.967287  2.337389  1.799443  1.135685\n",
       "333  2.974430  3.537706  1.987065 -1.861848  0.262459\n",
       "211  1.353304 -3.212935  1.034578  1.797846  1.682174\n",
       "850  1.205781  2.802149  4.308715 -1.498583  2.202188\n",
       "514  1.659687 -2.176888  2.090389  1.608092  1.805161\n",
       "..        ...       ...       ...       ...       ...\n",
       "95        NaN  3.770849  1.538188 -3.086773 -0.145386\n",
       "317  2.553215  3.059551  2.469591 -2.561048  2.554759\n",
       "481  1.915074 -2.785880  2.879047 -0.472673  5.318281\n",
       "140  2.302969  3.789408  0.752376 -2.758305  1.449600\n",
       "291 -0.034102 -0.700874  0.013272  2.687944  4.080553\n",
       "\n",
       "[800 rows x 5 columns]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = create_df()\n",
    "X_train, X_test, y_train, y_test = split_data(df, label=\"label\")\n",
    "X_train"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 1\n",
    "\n",
    "![scenario_1](./imgs/scenario_1.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[9.14049554e-06 9.99990860e-01]\n",
      " [8.96087282e-06 9.99991039e-01]\n",
      " [1.27023722e-05 9.99987298e-01]\n",
      " [9.14322593e-06 9.99990857e-01]\n",
      " [9.99966776e-01 3.32243731e-05]]\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline([\n",
    "    (\"imputer\", dp.BasicImputer(verbose=False)),\n",
    "    (\"encoder\", dp.DataMinMaxScaler(verbose=False)),\n",
    "    (\"model\", get_model())\n",
    "])\n",
    "pipe.fit(X_train, y_train)\n",
    "predictions = pipe.predict_proba(X_test)\n",
    "print(predictions[0:5])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 2\n",
    "\n",
    "![scenario_2](./imgs/scenario_2.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[8.22907996e-05 9.99917709e-01]\n",
      " [8.99293079e-06 9.99991007e-01]\n",
      " [4.20003325e-05 9.99958000e-01]\n",
      " [9.14326620e-06 9.99990857e-01]\n",
      " [9.99967810e-01 3.21900476e-05]]\n"
     ]
    }
   ],
   "source": [
    "# Define the pipeline to be used for each cohort\n",
    "cohort_pipeline = [\n",
    "    dp.BasicImputer(verbose=False),\n",
    "    dp.DataMinMaxScaler(verbose=False),\n",
    "]\n",
    "# Define the cohorts\n",
    "c1 = [ ['num_0', '==', np.nan], 'or', ['num_0', '>', 1.7] ]\n",
    "c2 = None\n",
    "# Create the cohort manager\n",
    "cohort_set = CohortManager(\n",
    "    transform_pipe=cohort_pipeline,\n",
    "    cohort_def=[c1, c2]\n",
    ")\n",
    "# Create a pipeline that uses the cohort manager\n",
    "pipe = Pipeline([\n",
    "    (\"cohort_preprocess\", cohort_set),\n",
    "    (\"model\", get_model())\n",
    "])\n",
    "pipe.fit(X_train, y_train)\n",
    "predictions = pipe.predict_proba(X_test)\n",
    "print(predictions[0:5])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 3\n",
    "\n",
    "![scenario_3](./imgs/scenario_3.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[9.14325263e-06 9.99990857e-01]\n",
      " [9.12541931e-06 9.99990875e-01]\n",
      " [3.26334434e-04 9.99673666e-01]\n",
      " [9.14325446e-06 9.99990857e-01]\n",
      " [9.99966824e-01 3.31763322e-05]]\n"
     ]
    }
   ],
   "source": [
    "# Define the cohorts\n",
    "c1 = [ ['num_0', '==', np.nan], 'or', ['num_0', '>', 1.7] ]\n",
    "c2 = None\n",
    "# Create each cohort's pipeline\n",
    "c1_pipe = [dp.BasicImputer(verbose=False)]\n",
    "c2_pipe = [dp.DataMinMaxScaler(verbose=False)]\n",
    "# Create the cohort manager\n",
    "cohort_set = CohortManager(\n",
    "    transform_pipe=[c1_pipe, c2_pipe],\n",
    "    cohort_def={\"high_num_0\":c1, \"low_num_0\":c2}\n",
    ")\n",
    "# Create a pipeline that uses the cohort manager\n",
    "pipe = Pipeline([\n",
    "    (\"cohort_preprocess\", cohort_set),\n",
    "    (\"model\", get_model())\n",
    "])\n",
    "pipe.fit(X_train, y_train)\n",
    "predictions = pipe.predict_proba(X_test)\n",
    "print(predictions[0:5])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 4\n",
    "\n",
    "![scenario_4](./imgs/scenario_4.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1.38884203e-05 9.99986112e-01]\n",
      " [7.43531292e-06 9.99992565e-01]\n",
      " [2.95033437e-05 9.99970497e-01]\n",
      " [7.64702060e-06 9.99992353e-01]\n",
      " [9.99960620e-01 3.93798545e-05]]\n"
     ]
    }
   ],
   "source": [
    "# Define the cohorts\n",
    "c1 = [ ['num_0', '>', 0.5] ]\n",
    "c2 = None\n",
    "# Create the cohort manager\n",
    "cohort_set = CohortManager(\n",
    "    transform_pipe=[get_model()],\n",
    "    cohort_def={\"high_num_0\":c1, \"low_num_0\":c2}\n",
    ")\n",
    "# Create a pipeline that uses the cohort manager\n",
    "pipe = Pipeline([\n",
    "    (\"imputer\", dp.BasicImputer(verbose=False)),\n",
    "    (\"encoder\", dp.DataMinMaxScaler(verbose=False)),\n",
    "    (\"model_cohort\", cohort_set)\n",
    "])\n",
    "pipe.fit(X_train, y_train)\n",
    "predictions = pipe.predict_proba(X_test)\n",
    "print(predictions[0:5])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 5\n",
    "\n",
    "![scenario_5](./imgs/scenario_5.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[8.18261911e-06 9.99991817e-01]\n",
      " [1.27965618e-05 9.99987203e-01]\n",
      " [1.10660836e-05 9.99988934e-01]\n",
      " [8.10620716e-06 9.99991894e-01]\n",
      " [9.99964873e-01 3.51269116e-05]]\n",
      "\n",
      "['cohort_0', 'cohort_1']\n",
      "cohort_0 pred shape: (134, 2)\n",
      "cohort_1 pred shape: (66, 2)\n"
     ]
    }
   ],
   "source": [
    "# Define the pipeline to be used for each cohort\n",
    "cohort_pipeline = [\n",
    "    dp.BasicImputer(verbose=False),\n",
    "    dp.DataMinMaxScaler(verbose=False),\n",
    "    get_model()\n",
    "]\n",
    "# Define the cohorts\n",
    "c1 = [ ['num_0', '==', np.nan], 'or', ['num_0', '>', 1.7] ]\n",
    "c2 = None\n",
    "# Create the cohort manager\n",
    "cohort_set = CohortManager(\n",
    "    transform_pipe=cohort_pipeline,\n",
    "    cohort_def=[c1, c2]\n",
    ")\n",
    "# Fit and predict using the CohortManager\n",
    "cohort_set.fit(X_train, y_train)\n",
    "predictions = cohort_set.predict_proba(X_test)\n",
    "print(f\"{predictions[0:5]}\\n\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since this time we're using the CohortManager directly, we can choose to split the predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = cohort_set.predict_proba(X_test, split_pred=True)\n",
    "print(f\"{list(predictions.keys())}\")\n",
    "print(f\"cohort_0 pred shape: {predictions['cohort_0'].shape}\")\n",
    "print(f\"cohort_1 pred shape: {predictions['cohort_1'].shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenario 6\n",
    "\n",
    "![scenario_6](./imgs/scenario_6.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1.05048860e-05 9.99989495e-01]\n",
      " [8.01234692e-06 9.99991988e-01]\n",
      " [1.10660836e-05 9.99988934e-01]\n",
      " [8.10747765e-06 9.99991893e-01]\n",
      " [9.99957260e-01 4.27396398e-05]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Define the cohorts\n",
    "c1 = [ ['num_0', '==', np.nan], 'or', ['num_0', '>', 1.7] ]\n",
    "c2 = None\n",
    "# Create each cohort's pipeline\n",
    "c1_pipe = [dp.BasicImputer(verbose=False), get_model()]\n",
    "c2_pipe = [dp.DataMinMaxScaler(verbose=False), get_model()]\n",
    "# Create the cohort manager\n",
    "cohort_set = CohortManager(\n",
    "    transform_pipe=[c1_pipe, c2_pipe],\n",
    "    cohort_def={\"high_num_0\":c1, \"low_num_0\":c2}\n",
    ")\n",
    "# Fit and predict using the CohortManager\n",
    "cohort_set.fit(X_train, y_train)\n",
    "predictions = cohort_set.predict_proba(X_test)\n",
    "print(f\"{predictions[0:5]}\\n\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once again, since we're using the ``predict()`` method of the CohortManager, we can split the predictions, which results in a dictionary with the predictions for each cohort."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['high_num_0', 'low_num_0']\n",
      "high_num_0 pred shape: (134, 2)\n",
      "low_num_0 pred shape: (66, 2)\n"
     ]
    }
   ],
   "source": [
    "predictions = cohort_set.predict_proba(X_test, split_pred=True)\n",
    "print(f\"{list(predictions.keys())}\")\n",
    "print(f\"high_num_0 pred shape: {predictions['high_num_0'].shape}\")\n",
    "print(f\"low_num_0 pred shape: {predictions['low_num_0'].shape}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "raipub",
   "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.9.15 (main, Nov 24 2022, 14:31:59) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "98cf402749abf383affb54f23cdde06b52ae2a6e4394659b91d1cafca4224ff4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}