notebooks/cohort/case_study/case_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cohort Case Study 1\n",
    "\n",
    "In this first case study for the ``cohort`` submodule, we'll create a synthetic dataset that has some characteristics that will highlight the advantages of using pre-processing pipelines for each cohort separately instead of using it over the entire dataset.\n",
    "\n",
    "Adopting a separate pipeline for each cohort is not an approach that will work in every situation. In fact, it is usually recommended to process the entire dataset instead of analyzing each cohort separately. But in some cases, certain cohorts behave differently from others, and the intensity of this difference will indicate if using a cohort-based pipeline is the best approach or not. When each cohort has a distinct behavior, *e.g.*, considerably different class distributions, different feature importance, opposite class behaviors, different value distribution for certain features, etc., then one cohort might end up degrading the performance of a model trained over the entire dataset when evaluated over the remaining cohorts. This is especially true when we have cohorts with very different behaviors, and one cohort that comprises the majority of the instances: in this case, the model trained over the entire dataset might simply learn how to predict the class of the instances of this majority cohort, and simply ignore the other cohorts. By doing this, the model will still achieve good results, but at the cost of neglecting the minority cohorts. This becomes a major concern when we are dealing with sensitive cohorts, that is, cohorts built with sensitive features, such as (but not limited to): gender, nationality, race, age, or a combination of these features. When a model performs well for one of the sensitive cohorts, but under-performs for the remaining cohorts, then the model is considered to be biased and may result in several legal problems. To mitigate these discrepancies, we might need to apply different pre-processing operations over each cohort separately, or even train different models for each cohort. \n",
    "\n",
    "\n",
    "## Creating the artificial dataset\n",
    "\n",
    "Given that scenarios that benefit from using different data pipelines for each cohort are not common, in this first case study we'll use a synthetic dataset that is created to artificially create this scenario. Here are the main characteristics that we want to see in our dataset:\n",
    "\n",
    "1. **Different cohorts and sub-cohorts that behave differently:** to simulate this, we can use different rules to establish each instance's class based on which cohort it belongs to. By doing this, a model will have a hard time understanding the general classification rule if trained with the entire dataset, since learning how to classify instances from one cohort may harm the classification capabilities for another cohort;\n",
    "2. **Different value ranges and distributions for numerical features:** if the values for a feature vary considerably for different cohorts, we may find it useful to normalize these features separately for each cohort, instead of normalizing it for the entire dataset.\n",
    "\n",
    "Our dataset will detail if a company went bankrupt (after a fixed number of months) or not. The only features used for each company are: the company's country's of origin, the industry sector to which it belongs, and the initial investment poured into the creation of the company measured in the company's country local currency. Each country has different local currency values, so it is expected that this feature varies based on the country's value. Also, due to each country's many characteristics (culture, environment, financial situation, social differences, etc.), each industry sector functions differently based on the country of the company. For example, a company that sells tropical fruits requires a lot less investment to succeed in tropical countries when compared to countries farther away from tropical areas. Finally, the rule adopted to define the class of each company (1 if the company went bankrupt, and 0 otherwise) is the following: if  the investment value is larger than a given threshold then the company succeeded (class 0). For some companies, however, this behavior is inverted: if the investment is larger than the threshold, the company goes bankrupt. This threshold is defined for each sector in each country.\n",
    "\n",
    "We'll also add some noise to the dataset by adding some missing values in the ``investment`` column, as well as inverting a small percentage of the final classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xgboost as xgb\n",
    "from sklearn.pipeline import Pipeline\n",
    "from lightgbm import LGBMClassifier\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from raimitigations.utils import split_data\n",
    "import raimitigations.dataprocessing as dp\n",
    "from raimitigations.cohort import (\n",
    "    CohortDefinition,\n",
    "    CohortManager,\n",
    "    fetch_cohort_results\n",
    ")\n",
    "\n",
    "SEED = 51\n",
    "#SEED = None\n",
    "np.random.seed(SEED)\n",
    "random.seed(SEED)\n",
    "\n",
    "def _create_country_df(samples: int, sectors: dict, country_name: str):\n",
    "    df = None\n",
    "    for key in sectors.keys():\n",
    "        size = int(samples * sectors[key][\"prob_occur\"])\n",
    "        invest = np.random.uniform(low=sectors[key][\"min\"], high=sectors[key][\"max\"], size=size)\n",
    "        min_invest = min(invest)\n",
    "        max_invest = max(invest)\n",
    "        range_invest = max_invest - min_invest\n",
    "        bankrupt_th = sectors[key][\"prob_success\"] * range_invest\n",
    "        inverted_behavior = sectors[key][\"inverted_behavior\"]\n",
    "        bankrupt = []\n",
    "        for i in range(invest.shape[0]):\n",
    "            inst_class = 1\n",
    "            if invest[i] > bankrupt_th:\n",
    "                inst_class = 0\n",
    "            if inverted_behavior:\n",
    "                inst_class = int(not inst_class)\n",
    "            bankrupt.append(inst_class)\n",
    "        noise_ind = np.random.choice(range(size), int(size*0.05), replace=False)\n",
    "        for ind in noise_ind:\n",
    "            bankrupt[ind] = int(not bankrupt[ind])\n",
    "        noise_ind = np.random.choice(range(size), int(size*0.1), replace=False)\n",
    "        for ind in noise_ind:\n",
    "            invest[ind] = np.nan\n",
    "        \n",
    "        country_col = [country_name for _ in range(size)]\n",
    "        sector_col = [key for _ in range(size)]\n",
    "        df_sector = pd.DataFrame({\n",
    "            \"investment\":invest,\n",
    "            \"sector\":sector_col,\n",
    "            \"country\":country_col,\n",
    "            \"bankrupt\":bankrupt\n",
    "        })\n",
    "        \n",
    "        if df is None:\n",
    "            df = df_sector\n",
    "        else:\n",
    "            df = pd.concat([df, df_sector], axis=0)\n",
    "    return df\n",
    "\n",
    "def create_df_multiple_distributions(samples: list):\n",
    "    sectors_c1 = {\n",
    "        \"s1\": {\"prob_occur\":0.5, \"prob_success\":0.99, \"inverted_behavior\":False, \"min\":2e6, \"max\":1e7},\n",
    "        \"s2\": {\"prob_occur\":0.1, \"prob_success\":0.2, \"inverted_behavior\":False, \"min\":1e7, \"max\":1.5e9},\n",
    "        \"s3\": {\"prob_occur\":0.1, \"prob_success\":0.9, \"inverted_behavior\":True, \"min\":1e9, \"max\":1e10},\n",
    "        \"s4\": {\"prob_occur\":0.3, \"prob_success\":0.7, \"inverted_behavior\":False, \"min\":4e10, \"max\":9e13},\n",
    "    }\n",
    "    sectors_c2 = {\n",
    "        \"s1\": {\"prob_occur\":0.1, \"prob_success\":0.6, \"inverted_behavior\":True, \"min\":1e3, \"max\":5e3},\n",
    "        \"s2\": {\"prob_occur\":0.3, \"prob_success\":0.9, \"inverted_behavior\":False, \"min\":1e5, \"max\":1.5e6},\n",
    "        \"s3\": {\"prob_occur\":0.5, \"prob_success\":0.3, \"inverted_behavior\":False, \"min\":5e4, \"max\":3e5},\n",
    "        \"s4\": {\"prob_occur\":0.1, \"prob_success\":0.8, \"inverted_behavior\":False, \"min\":1e6, \"max\":1e7},\n",
    "    }\n",
    "    sectors_c3 = {\n",
    "        \"s1\": {\"prob_occur\":0.3, \"prob_success\":0.9, \"inverted_behavior\":False, \"min\":3e2, \"max\":6e2},\n",
    "        \"s2\": {\"prob_occur\":0.6, \"prob_success\":0.7, \"inverted_behavior\":False, \"min\":5e3, \"max\":9e3},\n",
    "        \"s3\": {\"prob_occur\":0.07, \"prob_success\":0.6, \"inverted_behavior\":False, \"min\":4e3, \"max\":2e4},\n",
    "        \"s4\": {\"prob_occur\":0.03, \"prob_success\":0.1, \"inverted_behavior\":True, \"min\":6e5, \"max\":1.3e6},\n",
    "    }\n",
    "    countries = {\n",
    "        \"A\":{\"sectors\":sectors_c1, \"sample_rate\":0.85},\n",
    "        \"B\":{\"sectors\":sectors_c2, \"sample_rate\":0.05},\n",
    "        \"C\":{\"sectors\":sectors_c2, \"sample_rate\":0.1}\n",
    "    }\n",
    "    df = None\n",
    "    for key in countries.keys():\n",
    "        n_sample = int(samples * countries[key][\"sample_rate\"])\n",
    "        df_c = _create_country_df(n_sample, countries[key][\"sectors\"], key)\n",
    "        if df is None:\n",
    "            df = df_c\n",
    "        else:\n",
    "            df = pd.concat([df, df_c], axis=0)\n",
    "    \n",
    "    idx = pd.Index([i for i in range(df.shape[0])])\n",
    "    df = df.set_index(idx)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now create our artificial dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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>investment</th>\n",
       "      <th>sector</th>\n",
       "      <th>country</th>\n",
       "      <th>bankrupt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7.405851e+06</td>\n",
       "      <td>s1</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.357697e+06</td>\n",
       "      <td>s1</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4.746429e+06</td>\n",
       "      <td>s1</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7.152158e+06</td>\n",
       "      <td>s1</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>NaN</td>\n",
       "      <td>s1</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9995</th>\n",
       "      <td>4.226512e+06</td>\n",
       "      <td>s4</td>\n",
       "      <td>C</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9996</th>\n",
       "      <td>3.566758e+06</td>\n",
       "      <td>s4</td>\n",
       "      <td>C</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9997</th>\n",
       "      <td>9.281006e+06</td>\n",
       "      <td>s4</td>\n",
       "      <td>C</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9998</th>\n",
       "      <td>5.770378e+06</td>\n",
       "      <td>s4</td>\n",
       "      <td>C</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9999</th>\n",
       "      <td>3.661511e+06</td>\n",
       "      <td>s4</td>\n",
       "      <td>C</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10000 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "        investment sector country  bankrupt\n",
       "0     7.405851e+06     s1       A         1\n",
       "1     2.357697e+06     s1       A         1\n",
       "2     4.746429e+06     s1       A         1\n",
       "3     7.152158e+06     s1       A         1\n",
       "4              NaN     s1       A         1\n",
       "...            ...    ...     ...       ...\n",
       "9995  4.226512e+06     s4       C         1\n",
       "9996  3.566758e+06     s4       C         0\n",
       "9997  9.281006e+06     s4       C         0\n",
       "9998  5.770378e+06     s4       C         1\n",
       "9999  3.661511e+06     s4       C         1\n",
       "\n",
       "[10000 rows x 4 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = create_df_multiple_distributions(10000)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll now split our dataset into train and test sets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = split_data(df, label=\"bankrupt\", test_size=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    #model = LGBMClassifier(random_state=SEED)\n",
    "    model = LogisticRegression(random_state=SEED)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyzing the \"country\" cohorts\n",
    "\n",
    "Let's create our baseline model. We'll use a simple model since our goal is to test the efficiency of data processing pipelines, not test how different models behave. We'll create a pipeline with an imputer, a data normalization transformer, a one-hot encoding transformer, and finally our simple estimator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#EXPERIMENT: Baseline 1\n",
    "\n",
    "pipe = Pipeline([\n",
    "            (\"imputer\", dp.BasicImputer(verbose=False)),\n",
    "            (\"scaler\", dp.DataMinMaxScaler(verbose=False)),\n",
    "            (\"encoder\", dp.EncoderOHE(verbose=False)),\n",
    "            (\"estimator\", get_model()),\n",
    "        ])\n",
    "pipe.fit(X_train, y_train)\n",
    "pred_org = pipe.predict_proba(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll then test our pipeline over the test set and analyze how this pipeline performs over different cohorts. This analysis is done using the ``fetch_cohort_results`` function, which shows the results obtained for the entire dataset, as well as for different cohorts. Notice that this function computes all metrics for each cohort separately, and therefore, different thresholds might be encountered for each cohort (the optimal threshold of a given set of predictions is found using the ROC curve, and this threshold is used to determine if a probability should be converted to class ``1`` or class ``0``). Since we use the thresholds for computing the precision, recall, accuracy and F1 score, we must compute these thresholds using the training set, otherwise, we'll have some data leakage. Therefore, we first call the ``fetch_cohort_results`` using the training set. Notice that we set the ``return_th_dict`` to True (its default value is False), which makes the function return a dataframe with all of the metrics computed and a dictionary with the best thresholds found for each cohort. We then use this dict of thresholds when we call the ``fetch_cohort_results`` for a second time, but this time for the test dataset. Since we want to use the thresholds computed for the training set, we set the ``fixed_th`` parameter to be the dict of thresholds.\n",
    "\n",
    "Notice that the thresholds are only used for binary problems. For multi-class problems, the class chosen based on the probabilities is the class with the largest probability.\n",
    "\n",
    "For now, we'll focus on the cohorts defined by the different countries in the dataset.\n",
    "\n",
    "<a id='baseline1'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.803805</td>\n",
       "      <td>0.788653</td>\n",
       "      <td>0.790248</td>\n",
       "      <td>0.768984</td>\n",
       "      <td>0.769000</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.836172</td>\n",
       "      <td>0.821859</td>\n",
       "      <td>0.832952</td>\n",
       "      <td>0.813529</td>\n",
       "      <td>0.814303</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.788575</td>\n",
       "      <td>0.797303</td>\n",
       "      <td>0.784403</td>\n",
       "      <td>0.786350</td>\n",
       "      <td>0.790960</td>\n",
       "      <td>0.176188</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.798822</td>\n",
       "      <td>0.846590</td>\n",
       "      <td>0.830079</td>\n",
       "      <td>0.829746</td>\n",
       "      <td>0.832192</td>\n",
       "      <td>0.234479</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.803805   0.788653  0.790248  0.768984   \n",
       "1  cohort_0  (`country` == \"A\")  0.836172   0.821859  0.832952  0.813529   \n",
       "2  cohort_1  (`country` == \"B\")  0.788575   0.797303  0.784403  0.786350   \n",
       "3  cohort_2  (`country` == \"C\")  0.798822   0.846590  0.830079  0.829746   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.769000   0.714791      3000  \n",
       "1  0.814303   0.714791      2531  \n",
       "2  0.790960   0.176188       177  \n",
       "3  0.832192   0.234479       292  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_train_org = pipe.predict_proba(X_train)\n",
    "metrics_train, th_dict = fetch_cohort_results(X_train, y_train, pred_train_org, cohort_col=[\"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_org, cohort_col=[\"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that our pipeline managed to get a decent performance for all countries, despite the different behaviors that we injected for each cohort. However, since we have only a single estimator, we should consider only a single threshold for the entire dataset. The function ``fetch_cohort_results()`` analyzes the results separately for each cohort, that is, all metrics are computed using the isolated predictions of each cohort. That is why we can see that each cohort used different thresholds (``threshold`` column). But if we use the same optimal threshold computed for the entire dataset for all cohorts (which makes sense in this case, since we have a single estimator), then we'll notice that the results are very different.\n",
    "\n",
    "In the following cell, we call the ``fetch_cohort_results()`` again using the test dataset, but this time we set the ``shared_th`` to True, and we also specify that the thresholds to be used are the ones from ``th_dict`` (which were computed using the training set). This way, what will happen is that the precision, recall, accuracy, and F1 score metrics will be computed for all cohorts using the threshold value for the \"all\" cohort (because ``shared_th`` is True).\n",
    "\n",
    "<a id='baseline2'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.803805</td>\n",
       "      <td>0.788653</td>\n",
       "      <td>0.790248</td>\n",
       "      <td>0.768984</td>\n",
       "      <td>0.769000</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.836172</td>\n",
       "      <td>0.821859</td>\n",
       "      <td>0.832952</td>\n",
       "      <td>0.813529</td>\n",
       "      <td>0.814303</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.788575</td>\n",
       "      <td>0.579861</td>\n",
       "      <td>0.548716</td>\n",
       "      <td>0.486034</td>\n",
       "      <td>0.525424</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.798822</td>\n",
       "      <td>0.550498</td>\n",
       "      <td>0.530459</td>\n",
       "      <td>0.475221</td>\n",
       "      <td>0.523973</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.803805   0.788653  0.790248  0.768984   \n",
       "1  cohort_0  (`country` == \"A\")  0.836172   0.821859  0.832952  0.813529   \n",
       "2  cohort_1  (`country` == \"B\")  0.788575   0.579861  0.548716  0.486034   \n",
       "3  cohort_2  (`country` == \"C\")  0.798822   0.550498  0.530459  0.475221   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.769000   0.714791      3000  \n",
       "1  0.814303   0.714791      2531  \n",
       "2  0.525424   0.714791       177  \n",
       "3  0.523973   0.714791       292  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Baseline 2\n",
    "fetch_cohort_results(X_test, y_test, pred_org, cohort_col=[\"country\"], shared_th=True, fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the sake of comparability, we'll use different thresholds for each cohort for all of our experiments from this point on (that is, ``shared_th`` will be set to False). This will make our analysis more straightforward and easier to understand.\n",
    "\n",
    "Let's see if we can improve these metrics (compared to [Baseline 1](#baseline1)) by applying some pre-processing steps over each cohort separately. To that end, let's apply the imputation and normalization over each cohort separately and see how this impacts the resulting pipeline. We'll use the ``CohortManager`` class to achieve this.\n",
    "\n",
    "<a id='cohort1'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.814191</td>\n",
       "      <td>0.795227</td>\n",
       "      <td>0.797164</td>\n",
       "      <td>0.775978</td>\n",
       "      <td>0.776000</td>\n",
       "      <td>0.717450</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.840887</td>\n",
       "      <td>0.826499</td>\n",
       "      <td>0.838250</td>\n",
       "      <td>0.819723</td>\n",
       "      <td>0.820624</td>\n",
       "      <td>0.717450</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.803594</td>\n",
       "      <td>0.801724</td>\n",
       "      <td>0.803273</td>\n",
       "      <td>0.801855</td>\n",
       "      <td>0.802260</td>\n",
       "      <td>0.367261</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.820083</td>\n",
       "      <td>0.846590</td>\n",
       "      <td>0.830079</td>\n",
       "      <td>0.829746</td>\n",
       "      <td>0.832192</td>\n",
       "      <td>0.265146</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.814191   0.795227  0.797164  0.775978   \n",
       "1  cohort_0  (`country` == \"A\")  0.840887   0.826499  0.838250  0.819723   \n",
       "2  cohort_1  (`country` == \"B\")  0.803594   0.801724  0.803273  0.801855   \n",
       "3  cohort_2  (`country` == \"C\")  0.820083   0.846590  0.830079  0.829746   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.776000   0.717450      3000  \n",
       "1  0.820624   0.717450      2531  \n",
       "2  0.802260   0.367261       177  \n",
       "3  0.832192   0.265146       292  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Cohort 1\n",
    "\n",
    "cht_manager = CohortManager(\n",
    "    transform_pipe=[\n",
    "        dp.BasicImputer(verbose=False),\n",
    "        dp.DataMinMaxScaler(verbose=False),\n",
    "    ],\n",
    "    cohort_col=[\"country\"]\n",
    ")\n",
    "\n",
    "pipe = Pipeline([\n",
    "            (\"cht_preprocess\", cht_manager),\n",
    "            (\"encoder\", dp.EncoderOHE(verbose=False)),\n",
    "            (\"estimator\", get_model()),\n",
    "        ])\n",
    "pipe.fit(X_train, y_train)\n",
    "pred_cht = pipe.predict_proba(X_test)\n",
    "\n",
    "pred_train = pipe.predict_proba(X_train)\n",
    "metrics_train, th_dict = fetch_cohort_results(X_train, y_train, pred_train, cohort_col=[\"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_cht, cohort_col=[\"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, we only achieved a slight performance increase. Let's now try using the same pipeline used for the baseline model, but this time each cohort has its own pipeline, that is, the pre-processing steps and the estimator are fitted for each cohort separately.\n",
    "\n",
    "<a id='cohort'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.829600</td>\n",
       "      <td>0.813206</td>\n",
       "      <td>0.818502</td>\n",
       "      <td>0.799828</td>\n",
       "      <td>0.800000</td>\n",
       "      <td>0.725327</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.837492</td>\n",
       "      <td>0.827085</td>\n",
       "      <td>0.838913</td>\n",
       "      <td>0.820496</td>\n",
       "      <td>0.821414</td>\n",
       "      <td>0.725327</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.830680</td>\n",
       "      <td>0.797303</td>\n",
       "      <td>0.784403</td>\n",
       "      <td>0.786350</td>\n",
       "      <td>0.790960</td>\n",
       "      <td>0.110570</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.889919</td>\n",
       "      <td>0.846590</td>\n",
       "      <td>0.830079</td>\n",
       "      <td>0.829746</td>\n",
       "      <td>0.832192</td>\n",
       "      <td>0.174096</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.829600   0.813206  0.818502  0.799828   \n",
       "1  cohort_0  (`country` == \"A\")  0.837492   0.827085  0.838913  0.820496   \n",
       "2  cohort_1  (`country` == \"B\")  0.830680   0.797303  0.784403  0.786350   \n",
       "3  cohort_2  (`country` == \"C\")  0.889919   0.846590  0.830079  0.829746   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.800000   0.725327      3000  \n",
       "1  0.821414   0.725327      2531  \n",
       "2  0.790960   0.110570       177  \n",
       "3  0.832192   0.174096       292  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Cohort 2\n",
    "\n",
    "cht_manager = CohortManager(\n",
    "    transform_pipe=[\n",
    "        dp.BasicImputer(verbose=False),\n",
    "        dp.DataMinMaxScaler(verbose=False),\n",
    "        dp.EncoderOHE(verbose=False),\n",
    "        get_model()\n",
    "    ],\n",
    "    cohort_col=[\"country\"]\n",
    ")\n",
    "cht_manager.fit(X_train, y_train)\n",
    "pred_cht = cht_manager.predict_proba(X_test)\n",
    "\n",
    "pred_train = cht_manager.predict_proba(X_train)\n",
    "metrics_train, th_dict = fetch_cohort_results(X_train, y_train, pred_train, cohort_col=[\"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_cht, cohort_col=[\"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We managed to get a decent increase in the metrics by simply training different models for each cohort. In this case, we trained different pipelines for each country cohort. However, each industry sector behaves differently, so even though we are now looking at each country separately, the ``sector`` column still hinders the trained model. Let's now replicate the previous experiment, but this time train different pipelines for each industry sector instead of each country:\n",
    "\n",
    "<a id='cohort3'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.902798</td>\n",
       "      <td>0.902903</td>\n",
       "      <td>0.895324</td>\n",
       "      <td>0.898555</td>\n",
       "      <td>0.902333</td>\n",
       "      <td>0.464561</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.924799</td>\n",
       "      <td>0.927963</td>\n",
       "      <td>0.917449</td>\n",
       "      <td>0.921941</td>\n",
       "      <td>0.925721</td>\n",
       "      <td>0.464561</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.747497</td>\n",
       "      <td>0.797229</td>\n",
       "      <td>0.796085</td>\n",
       "      <td>0.790934</td>\n",
       "      <td>0.790960</td>\n",
       "      <td>0.747629</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.808819</td>\n",
       "      <td>0.795423</td>\n",
       "      <td>0.793871</td>\n",
       "      <td>0.794047</td>\n",
       "      <td>0.794521</td>\n",
       "      <td>0.504412</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.902798   0.902903  0.895324  0.898555   \n",
       "1  cohort_0  (`country` == \"A\")  0.924799   0.927963  0.917449  0.921941   \n",
       "2  cohort_1  (`country` == \"B\")  0.747497   0.797229  0.796085  0.790934   \n",
       "3  cohort_2  (`country` == \"C\")  0.808819   0.795423  0.793871  0.794047   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.902333   0.464561      3000  \n",
       "1  0.925721   0.464561      2531  \n",
       "2  0.790960   0.747629       177  \n",
       "3  0.794521   0.504412       292  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Cohort 3\n",
    "\n",
    "cht_manager = CohortManager(\n",
    "    transform_pipe=[\n",
    "        dp.BasicImputer(verbose=False),\n",
    "        dp.DataMinMaxScaler(verbose=False),\n",
    "        dp.EncoderOHE(verbose=False),\n",
    "        get_model()\n",
    "    ],\n",
    "    cohort_col=[\"sector\"]\n",
    ")\n",
    "cht_manager.fit(X_train, y_train)\n",
    "pred_cht = cht_manager.predict_proba(X_test)\n",
    "\n",
    "pred_train = cht_manager.predict_proba(X_train)\n",
    "metrics_train, th_dict = fetch_cohort_results(X_train, y_train, pred_train, cohort_col=[\"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_cht, cohort_col=[\"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the previous results, we can see that the difference in behavior for each sector is greater than the differences imposed by the countries. By training separate pipelines for each sector, we managed to greatly improve the performance of our model.\n",
    "\n",
    "Let's now see if we can improve these results even further using a different pipeline for each combination of country and sector:\n",
    "\n",
    "<a id='cohort4'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.921679</td>\n",
       "      <td>0.908100</td>\n",
       "      <td>0.902547</td>\n",
       "      <td>0.905012</td>\n",
       "      <td>0.908333</td>\n",
       "      <td>0.475074</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`country` == \"A\")</td>\n",
       "      <td>0.925494</td>\n",
       "      <td>0.920170</td>\n",
       "      <td>0.911027</td>\n",
       "      <td>0.914990</td>\n",
       "      <td>0.919004</td>\n",
       "      <td>0.476344</td>\n",
       "      <td>2531</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`country` == \"B\")</td>\n",
       "      <td>0.869063</td>\n",
       "      <td>0.845424</td>\n",
       "      <td>0.845956</td>\n",
       "      <td>0.841803</td>\n",
       "      <td>0.841808</td>\n",
       "      <td>0.546313</td>\n",
       "      <td>177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`country` == \"C\")</td>\n",
       "      <td>0.923664</td>\n",
       "      <td>0.874875</td>\n",
       "      <td>0.868682</td>\n",
       "      <td>0.869120</td>\n",
       "      <td>0.869863</td>\n",
       "      <td>0.388728</td>\n",
       "      <td>292</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     cohort           cht_query       roc  precision    recall        f1  \\\n",
       "0       all                 all  0.921679   0.908100  0.902547  0.905012   \n",
       "1  cohort_0  (`country` == \"A\")  0.925494   0.920170  0.911027  0.914990   \n",
       "2  cohort_1  (`country` == \"B\")  0.869063   0.845424  0.845956  0.841803   \n",
       "3  cohort_2  (`country` == \"C\")  0.923664   0.874875  0.868682  0.869120   \n",
       "\n",
       "   accuracy  threshold  cht_size  \n",
       "0  0.908333   0.475074      3000  \n",
       "1  0.919004   0.476344      2531  \n",
       "2  0.841808   0.546313       177  \n",
       "3  0.869863   0.388728       292  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Cohort 4\n",
    "\n",
    "cht_manager = CohortManager(\n",
    "    transform_pipe=[\n",
    "        dp.BasicImputer(verbose=False),\n",
    "        dp.DataMinMaxScaler(verbose=False),\n",
    "        dp.EncoderOHE(verbose=False),\n",
    "        get_model()\n",
    "    ],\n",
    "    cohort_col=[\"sector\", \"country\"]\n",
    ")\n",
    "cht_manager.fit(X_train, y_train)\n",
    "pred_cht = cht_manager.predict_proba(X_test)\n",
    "\n",
    "pred_train = cht_manager.predict_proba(X_train)\n",
    "metrics_train, th_dict = fetch_cohort_results(X_train, y_train, pred_train, cohort_col=[\"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_cht, cohort_col=[\"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This resulted in only a slight increase in performance. These results show us what we wanted: for our artificial dataset, training a separate pipeline for each cohort of sector and country, we get the best results. This was already expected due to how we created this artificial dataset, where we defined different classification behaviors for each subset with different sector and country values.\n",
    "\n",
    "## Checking the \"sector\" + \"country\" cohorts\n",
    "\n",
    "Let's now look into the metrics for each combination of country and sector. We'll first check how the baseline pipeline, trained over the entire dataset ([Baseline 1](#baseline1)), performs for each of these cohorts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.803805</td>\n",
       "      <td>0.788653</td>\n",
       "      <td>0.790248</td>\n",
       "      <td>0.768984</td>\n",
       "      <td>0.769000</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"A\")</td>\n",
       "      <td>0.852323</td>\n",
       "      <td>0.860370</td>\n",
       "      <td>0.896209</td>\n",
       "      <td>0.873737</td>\n",
       "      <td>0.888436</td>\n",
       "      <td>0.714791</td>\n",
       "      <td>1228</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"B\")</td>\n",
       "      <td>0.238095</td>\n",
       "      <td>0.208333</td>\n",
       "      <td>0.416667</td>\n",
       "      <td>0.277778</td>\n",
       "      <td>0.384615</td>\n",
       "      <td>0.732667</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"C\")</td>\n",
       "      <td>0.161905</td>\n",
       "      <td>0.232143</td>\n",
       "      <td>0.464286</td>\n",
       "      <td>0.309524</td>\n",
       "      <td>0.448276</td>\n",
       "      <td>0.796960</td>\n",
       "      <td>29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>cohort_3</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"A\")</td>\n",
       "      <td>0.803265</td>\n",
       "      <td>0.867622</td>\n",
       "      <td>0.834155</td>\n",
       "      <td>0.848937</td>\n",
       "      <td>0.893471</td>\n",
       "      <td>0.400690</td>\n",
       "      <td>291</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>cohort_4</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"B\")</td>\n",
       "      <td>0.813390</td>\n",
       "      <td>0.803529</td>\n",
       "      <td>0.867521</td>\n",
       "      <td>0.828205</td>\n",
       "      <td>0.880597</td>\n",
       "      <td>0.422347</td>\n",
       "      <td>67</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>cohort_5</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"C\")</td>\n",
       "      <td>0.775021</td>\n",
       "      <td>0.701744</td>\n",
       "      <td>0.787014</td>\n",
       "      <td>0.730828</td>\n",
       "      <td>0.858491</td>\n",
       "      <td>0.511512</td>\n",
       "      <td>106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>cohort_6</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"A\")</td>\n",
       "      <td>0.188799</td>\n",
       "      <td>0.575122</td>\n",
       "      <td>0.529032</td>\n",
       "      <td>0.285017</td>\n",
       "      <td>0.304878</td>\n",
       "      <td>0.163522</td>\n",
       "      <td>246</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>cohort_7</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"B\")</td>\n",
       "      <td>0.710065</td>\n",
       "      <td>0.913043</td>\n",
       "      <td>0.684211</td>\n",
       "      <td>0.721612</td>\n",
       "      <td>0.842105</td>\n",
       "      <td>0.176188</td>\n",
       "      <td>76</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>cohort_8</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"C\")</td>\n",
       "      <td>0.766817</td>\n",
       "      <td>0.889474</td>\n",
       "      <td>0.773175</td>\n",
       "      <td>0.814833</td>\n",
       "      <td>0.906977</td>\n",
       "      <td>0.234479</td>\n",
       "      <td>129</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>cohort_9</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"A\")</td>\n",
       "      <td>0.870353</td>\n",
       "      <td>0.935381</td>\n",
       "      <td>0.894539</td>\n",
       "      <td>0.911125</td>\n",
       "      <td>0.925587</td>\n",
       "      <td>0.503846</td>\n",
       "      <td>766</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>cohort_10</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"B\")</td>\n",
       "      <td>0.944444</td>\n",
       "      <td>0.875000</td>\n",
       "      <td>0.933333</td>\n",
       "      <td>0.892857</td>\n",
       "      <td>0.904762</td>\n",
       "      <td>0.944555</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>cohort_11</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"C\")</td>\n",
       "      <td>0.794444</td>\n",
       "      <td>0.804813</td>\n",
       "      <td>0.816667</td>\n",
       "      <td>0.809524</td>\n",
       "      <td>0.821429</td>\n",
       "      <td>0.960627</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       cohort                                  cht_query       roc  precision  \\\n",
       "0         all                                        all  0.803805   0.788653   \n",
       "1    cohort_0  (`sector` == \"s1\") and (`country` == \"A\")  0.852323   0.860370   \n",
       "2    cohort_1  (`sector` == \"s1\") and (`country` == \"B\")  0.238095   0.208333   \n",
       "3    cohort_2  (`sector` == \"s1\") and (`country` == \"C\")  0.161905   0.232143   \n",
       "4    cohort_3  (`sector` == \"s2\") and (`country` == \"A\")  0.803265   0.867622   \n",
       "5    cohort_4  (`sector` == \"s2\") and (`country` == \"B\")  0.813390   0.803529   \n",
       "6    cohort_5  (`sector` == \"s2\") and (`country` == \"C\")  0.775021   0.701744   \n",
       "7    cohort_6  (`sector` == \"s3\") and (`country` == \"A\")  0.188799   0.575122   \n",
       "8    cohort_7  (`sector` == \"s3\") and (`country` == \"B\")  0.710065   0.913043   \n",
       "9    cohort_8  (`sector` == \"s3\") and (`country` == \"C\")  0.766817   0.889474   \n",
       "10   cohort_9  (`sector` == \"s4\") and (`country` == \"A\")  0.870353   0.935381   \n",
       "11  cohort_10  (`sector` == \"s4\") and (`country` == \"B\")  0.944444   0.875000   \n",
       "12  cohort_11  (`sector` == \"s4\") and (`country` == \"C\")  0.794444   0.804813   \n",
       "\n",
       "      recall        f1  accuracy  threshold  cht_size  \n",
       "0   0.790248  0.768984  0.769000   0.714791      3000  \n",
       "1   0.896209  0.873737  0.888436   0.714791      1228  \n",
       "2   0.416667  0.277778  0.384615   0.732667        13  \n",
       "3   0.464286  0.309524  0.448276   0.796960        29  \n",
       "4   0.834155  0.848937  0.893471   0.400690       291  \n",
       "5   0.867521  0.828205  0.880597   0.422347        67  \n",
       "6   0.787014  0.730828  0.858491   0.511512       106  \n",
       "7   0.529032  0.285017  0.304878   0.163522       246  \n",
       "8   0.684211  0.721612  0.842105   0.176188        76  \n",
       "9   0.773175  0.814833  0.906977   0.234479       129  \n",
       "10  0.894539  0.911125  0.925587   0.503846       766  \n",
       "11  0.933333  0.892857  0.904762   0.944555        21  \n",
       "12  0.816667  0.809524  0.821429   0.960627        28  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Baseline 3\n",
    "\n",
    "_, th_dict = fetch_cohort_results(X_train, y_train, pred_train_org, cohort_col=[\"sector\", \"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_org, cohort_col=[\"sector\", \"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that some cohorts have very poor performance. For example, for companies from sector ``s1``, we can notice that the model learned how to predict the correct class only for companies from country ``A`` (``cohort_0``). Since companies from sector ``s1`` from country ``B`` used an inverted behavior (``inverted_behavior`` flag in the function that creates the dataset), and companies from country ``C`` had a very different range of investment values compared to the country ``A``, and since instances from country ``A`` are the majority of instances, then the model prioritized to learn how to classify only the instances of sector == ``s1`` from country ``A``.\n",
    "\n",
    "Let's now check the results for these cohorts using our best pipeline, which is the one where each cohort of sector and country had its own pipeline ([Cohort 4](#cohort4)):\n",
    "\n",
    "<a id='cohort5'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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>cohort</th>\n",
       "      <th>cht_query</th>\n",
       "      <th>roc</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>threshold</th>\n",
       "      <th>cht_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>all</td>\n",
       "      <td>all</td>\n",
       "      <td>0.921679</td>\n",
       "      <td>0.908100</td>\n",
       "      <td>0.902547</td>\n",
       "      <td>0.905012</td>\n",
       "      <td>0.908333</td>\n",
       "      <td>0.475074</td>\n",
       "      <td>3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>cohort_0</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"A\")</td>\n",
       "      <td>0.911655</td>\n",
       "      <td>0.941368</td>\n",
       "      <td>0.894939</td>\n",
       "      <td>0.914116</td>\n",
       "      <td>0.931596</td>\n",
       "      <td>0.445645</td>\n",
       "      <td>1228</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>cohort_1</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"B\")</td>\n",
       "      <td>0.904762</td>\n",
       "      <td>0.888889</td>\n",
       "      <td>0.833333</td>\n",
       "      <td>0.837500</td>\n",
       "      <td>0.846154</td>\n",
       "      <td>0.616965</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>cohort_2</td>\n",
       "      <td>(`sector` == \"s1\") and (`country` == \"C\")</td>\n",
       "      <td>0.914286</td>\n",
       "      <td>0.899038</td>\n",
       "      <td>0.895238</td>\n",
       "      <td>0.896057</td>\n",
       "      <td>0.896552</td>\n",
       "      <td>0.471948</td>\n",
       "      <td>29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>cohort_3</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"A\")</td>\n",
       "      <td>0.861716</td>\n",
       "      <td>0.867622</td>\n",
       "      <td>0.834155</td>\n",
       "      <td>0.848937</td>\n",
       "      <td>0.893471</td>\n",
       "      <td>0.488684</td>\n",
       "      <td>291</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>cohort_4</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"B\")</td>\n",
       "      <td>0.814815</td>\n",
       "      <td>0.914912</td>\n",
       "      <td>0.836895</td>\n",
       "      <td>0.868782</td>\n",
       "      <td>0.925373</td>\n",
       "      <td>0.521201</td>\n",
       "      <td>67</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>cohort_5</td>\n",
       "      <td>(`sector` == \"s2\") and (`country` == \"C\")</td>\n",
       "      <td>0.874276</td>\n",
       "      <td>0.929167</td>\n",
       "      <td>0.840778</td>\n",
       "      <td>0.878077</td>\n",
       "      <td>0.952830</td>\n",
       "      <td>0.662477</td>\n",
       "      <td>106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>cohort_6</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"A\")</td>\n",
       "      <td>0.875000</td>\n",
       "      <td>0.943813</td>\n",
       "      <td>0.877957</td>\n",
       "      <td>0.905093</td>\n",
       "      <td>0.934959</td>\n",
       "      <td>0.477076</td>\n",
       "      <td>246</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>cohort_7</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"B\")</td>\n",
       "      <td>0.762696</td>\n",
       "      <td>0.913043</td>\n",
       "      <td>0.684211</td>\n",
       "      <td>0.721612</td>\n",
       "      <td>0.842105</td>\n",
       "      <td>0.226486</td>\n",
       "      <td>76</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>cohort_8</td>\n",
       "      <td>(`sector` == \"s3\") and (`country` == \"C\")</td>\n",
       "      <td>0.780353</td>\n",
       "      <td>0.889474</td>\n",
       "      <td>0.773175</td>\n",
       "      <td>0.814833</td>\n",
       "      <td>0.906977</td>\n",
       "      <td>0.388728</td>\n",
       "      <td>129</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>cohort_9</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"A\")</td>\n",
       "      <td>0.907772</td>\n",
       "      <td>0.935381</td>\n",
       "      <td>0.894539</td>\n",
       "      <td>0.911125</td>\n",
       "      <td>0.925587</td>\n",
       "      <td>0.480231</td>\n",
       "      <td>766</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>cohort_10</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"B\")</td>\n",
       "      <td>0.900000</td>\n",
       "      <td>0.837500</td>\n",
       "      <td>0.800000</td>\n",
       "      <td>0.815249</td>\n",
       "      <td>0.857143</td>\n",
       "      <td>0.559088</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>cohort_11</td>\n",
       "      <td>(`sector` == \"s4\") and (`country` == \"C\")</td>\n",
       "      <td>0.833333</td>\n",
       "      <td>0.862500</td>\n",
       "      <td>0.822222</td>\n",
       "      <td>0.836257</td>\n",
       "      <td>0.857143</td>\n",
       "      <td>0.560495</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       cohort                                  cht_query       roc  precision  \\\n",
       "0         all                                        all  0.921679   0.908100   \n",
       "1    cohort_0  (`sector` == \"s1\") and (`country` == \"A\")  0.911655   0.941368   \n",
       "2    cohort_1  (`sector` == \"s1\") and (`country` == \"B\")  0.904762   0.888889   \n",
       "3    cohort_2  (`sector` == \"s1\") and (`country` == \"C\")  0.914286   0.899038   \n",
       "4    cohort_3  (`sector` == \"s2\") and (`country` == \"A\")  0.861716   0.867622   \n",
       "5    cohort_4  (`sector` == \"s2\") and (`country` == \"B\")  0.814815   0.914912   \n",
       "6    cohort_5  (`sector` == \"s2\") and (`country` == \"C\")  0.874276   0.929167   \n",
       "7    cohort_6  (`sector` == \"s3\") and (`country` == \"A\")  0.875000   0.943813   \n",
       "8    cohort_7  (`sector` == \"s3\") and (`country` == \"B\")  0.762696   0.913043   \n",
       "9    cohort_8  (`sector` == \"s3\") and (`country` == \"C\")  0.780353   0.889474   \n",
       "10   cohort_9  (`sector` == \"s4\") and (`country` == \"A\")  0.907772   0.935381   \n",
       "11  cohort_10  (`sector` == \"s4\") and (`country` == \"B\")  0.900000   0.837500   \n",
       "12  cohort_11  (`sector` == \"s4\") and (`country` == \"C\")  0.833333   0.862500   \n",
       "\n",
       "      recall        f1  accuracy  threshold  cht_size  \n",
       "0   0.902547  0.905012  0.908333   0.475074      3000  \n",
       "1   0.894939  0.914116  0.931596   0.445645      1228  \n",
       "2   0.833333  0.837500  0.846154   0.616965        13  \n",
       "3   0.895238  0.896057  0.896552   0.471948        29  \n",
       "4   0.834155  0.848937  0.893471   0.488684       291  \n",
       "5   0.836895  0.868782  0.925373   0.521201        67  \n",
       "6   0.840778  0.878077  0.952830   0.662477       106  \n",
       "7   0.877957  0.905093  0.934959   0.477076       246  \n",
       "8   0.684211  0.721612  0.842105   0.226486        76  \n",
       "9   0.773175  0.814833  0.906977   0.388728       129  \n",
       "10  0.894539  0.911125  0.925587   0.480231       766  \n",
       "11  0.800000  0.815249  0.857143   0.559088        21  \n",
       "12  0.822222  0.836257  0.857143   0.560495        28  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#EXPERIMENT: Cohort 5\n",
    "\n",
    "_, th_dict = fetch_cohort_results(X_train, y_train, pred_train, cohort_col=[\"sector\", \"country\"], return_th_dict=True)\n",
    "fetch_cohort_results(X_test, y_test, pred_cht, cohort_col=[\"sector\", \"country\"], fixed_th=th_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the results here are a lot more consistent, and there aren't any cohorts with a drastic performance difference when compared to other cohorts. This is because each cohort had its own estimators, so it had only one behavior to learn."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rai",
   "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": "5215fd1cce9b64233b490e6c4e868a314fa8a5a2b064fe0583fd0c29542af034"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}