fancy-imputation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fancy Imputation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fill nulls with group averages\n",
    "\n",
    "Suppose we have some data with missing values and we want to impute the missing values using the average of another group in the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df.shape: (5000, 2)\n"
     ]
    },
    {
     "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>flavor</th>\n",
       "      <th>pints</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>moolenium crunch</td>\n",
       "      <td>11.057570</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>6.288724</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>chubby hubby</td>\n",
       "      <td>7.660815</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>6.644338</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>neopolitan</td>\n",
       "      <td>13.600125</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             flavor      pints\n",
       "0  moolenium crunch  11.057570\n",
       "1         bubblegum   6.288724\n",
       "2      chubby hubby   7.660815\n",
       "3         bubblegum   6.644338\n",
       "4        neopolitan  13.600125"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('https://gist.githubusercontent.com/zgulde/0f2c364ea77075f0fdeef95eaf9b4a20/raw/0c9a72c26712822fa36320a469b432029d512e2c/ice_cream.csv')\n",
    "print('df.shape:', df.shape)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "flavor      0\n",
       "pints     982\n",
       "dtype: int64"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.isna().sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here some of the values for `pints` are null, but we have all the values for flavor. When doing imputation we might wish to fill the missing values not just with the overall average, but with the average for each flavor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Splitting\n",
    "\n",
    "Whatever parameters for imputation we calculate (e.g. the mean if we are filling missing values with the average) should be calculated from the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train shape:  (4050, 2) - no. of nulls in pints 788\n",
      "validate shape:  (450, 2) - no. of nulls in pints 86\n",
      "test shape:  (500, 2) - no. of nulls in pints 108\n"
     ]
    }
   ],
   "source": [
    "train_and_validate, test = train_test_split(df, test_size=.1, random_state=123)\n",
    "train, validate = train_test_split(train_and_validate, test_size=.1, random_state=123)\n",
    "\n",
    "print('train shape: ', train.shape, '- no. of nulls in pints', train.pints.isna().sum())\n",
    "print('validate shape: ', validate.shape, '- no. of nulls in pints', validate.pints.isna().sum())\n",
    "print('test shape: ', test.shape, '- no. of nulls in pints', test.pints.isna().sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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>flavor</th>\n",
       "      <th>pints</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>776</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1848</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1878</th>\n",
       "      <td>moolenium crunch</td>\n",
       "      <td>9.583060</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3475</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>5.638092</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4536</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>6.191578</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                flavor     pints\n",
       "776         rocky road       NaN\n",
       "1848         bubblegum       NaN\n",
       "1878  moolenium crunch  9.583060\n",
       "3475         bubblegum  5.638092\n",
       "4536         bubblegum  6.191578"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do this in one of two ways:\n",
    "\n",
    "1. Use `.replace` to fill missing values with the group averages.\n",
    "1. Join the table on to our dataframes using the group as the join key, then use the average column to fill missing values.\n",
    "\n",
    "The first option is much simpler, but will only work when we have a single group we are looking at. If we want to fill missing values based off of multiple groups, we will have to use the second option.\n",
    "\n",
    "We'll demonstrate both options here:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `.replace`\n",
    "\n",
    "To use `.replace`, we'll need a series of the average value for each group:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "flavor\n",
       "bubblegum            6.002209\n",
       "chubby hubby        10.054496\n",
       "moolenium crunch    10.056668\n",
       "neopolitan          10.061540\n",
       "pistachio           10.108655\n",
       "rocky road          15.135639\n",
       "Name: pints, dtype: float64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flavor_pint_avgs = train.groupby('flavor').pints.mean()\n",
    "flavor_pint_avgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that because we added `.pints` after the group by, we get a series as the output, not a dataframe.\n",
    "\n",
    "Now we can use `.replace`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "776     15.135639\n",
       "1848     6.002209\n",
       "1878    10.056668\n",
       "3475     6.002209\n",
       "4536     6.002209\n",
       "          ...    \n",
       "3343    10.056668\n",
       "4575    15.135639\n",
       "1008    10.054496\n",
       "1264    10.108655\n",
       "1634    10.056668\n",
       "Name: flavor, Length: 4050, dtype: float64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.flavor.replace(flavor_pint_avgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`.replace` will produce a series where each flavor has been replaced with the average number of pints for that flavor. That is, for each observation we have the average pint value for that observation's flavor.\n",
    "\n",
    "We can use this series to fill in the missing values for pints with the average value by flavor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.pints = train.pints.fillna(train.flavor.replace(flavor_pint_avgs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can do the same thing for the validate and test splits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "validate.pints = validate.pints.fillna(validate.flavor.replace(flavor_pint_avgs))\n",
    "test.pints = test.pints.fillna(test.flavor.replace(flavor_pint_avgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train shape:  (4050, 2) - no. of nulls in pints 0\n",
      "validate shape:  (450, 2) - no. of nulls in pints 0\n",
      "test shape:  (500, 2) - no. of nulls in pints 0\n"
     ]
    }
   ],
   "source": [
    "print('train shape: ', train.shape, '- no. of nulls in pints', train.pints.isna().sum())\n",
    "print('validate shape: ', validate.shape, '- no. of nulls in pints', validate.pints.isna().sum())\n",
    "print('test shape: ', test.shape, '- no. of nulls in pints', test.pints.isna().sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we haven't leaked any information from the validate or test splits because `flavor_pint_avgs` was calculated from the training data split."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Joins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train shape:  (4050, 2) - no. of nulls in pints 788\n",
      "validate shape:  (450, 2) - no. of nulls in pints 86\n",
      "test shape:  (500, 2) - no. of nulls in pints 108\n"
     ]
    }
   ],
   "source": [
    "train_and_validate, test = train_test_split(df, test_size=.1, random_state=123)\n",
    "train, validate = train_test_split(train_and_validate, test_size=.1, random_state=123)\n",
    "\n",
    "print('train shape: ', train.shape, '- no. of nulls in pints', train.pints.isna().sum())\n",
    "print('validate shape: ', validate.shape, '- no. of nulls in pints', validate.pints.isna().sum())\n",
    "print('test shape: ', test.shape, '- no. of nulls in pints', test.pints.isna().sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This method is a little more complicated, but will work when the value we want to impute with is calculated from multiple groups instead of just 1."
   ]
  },
  {
   "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>flavor</th>\n",
       "      <th>pints</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>bubblegum</td>\n",
       "      <td>6.002209</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>chubby hubby</td>\n",
       "      <td>10.054496</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>moolenium crunch</td>\n",
       "      <td>10.056668</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>neopolitan</td>\n",
       "      <td>10.061540</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             flavor      pints\n",
       "0         bubblegum   6.002209\n",
       "1      chubby hubby  10.054496\n",
       "2  moolenium crunch  10.056668\n",
       "3        neopolitan  10.061540\n",
       "4         pistachio  10.108655\n",
       "5        rocky road  15.135639"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flavor_pint_avgs = train.groupby(['flavor']).pints.mean().reset_index()\n",
    "flavor_pint_avgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice here the group averages table we are producing is a dataframe, not a series.\n",
    "\n",
    "Now we join back to the original dataframe:"
   ]
  },
  {
   "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>flavor</th>\n",
       "      <th>pints_x</th>\n",
       "      <th>pints_y</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>776</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>NaN</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1848</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>14.716520</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1878</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>13.448409</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3475</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>NaN</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4536</th>\n",
       "      <td>rocky road</td>\n",
       "      <td>13.674279</td>\n",
       "      <td>15.135639</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3343</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>11.688645</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4575</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>NaN</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1008</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>NaN</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1264</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>8.359803</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1634</th>\n",
       "      <td>pistachio</td>\n",
       "      <td>14.228036</td>\n",
       "      <td>10.108655</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4050 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          flavor    pints_x    pints_y\n",
       "776   rocky road        NaN  15.135639\n",
       "1848  rocky road  14.716520  15.135639\n",
       "1878  rocky road  13.448409  15.135639\n",
       "3475  rocky road        NaN  15.135639\n",
       "4536  rocky road  13.674279  15.135639\n",
       "...          ...        ...        ...\n",
       "3343   pistachio  11.688645  10.108655\n",
       "4575   pistachio        NaN  10.108655\n",
       "1008   pistachio        NaN  10.108655\n",
       "1264   pistachio   8.359803  10.108655\n",
       "1634   pistachio  14.228036  10.108655\n",
       "\n",
       "[4050 rows x 3 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx = train.index\n",
    "train_merged = train.merge(flavor_pint_avgs, on='flavor').set_index(idx)\n",
    "train_merged"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note several things here:\n",
    "\n",
    "- The `.merge` operation loses the index, so if we care about preserving the original indices, we need to store them in a variable ahead of time and re-set the index after the merge.\n",
    "- The resulting dataframe has two columns for pints, `pints_x` and `pints_y`.\n",
    "    - `pints_x` comes from the training dataset. This column still has null values.\n",
    "    - `pints_y` comes from the group averages dataframe. This column is the average pints value by flavor.\n",
    "    \n",
    "We want to create a single `pints` column that has all of the values from the training dataset, but where the value is null, we instead want the group average."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.pints = train_merged.pints_x.fillna(train_merged.pints_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can condense the code above and apply it to validate and test as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = validate.index\n",
    "validate_merged = validate.merge(flavor_pint_avgs, on='flavor').set_index(idx)\n",
    "validate.pints = validate_merged.pints_x.fillna(validate_merged.pints_y)\n",
    "\n",
    "idx = test.index\n",
    "test_merged = test.merge(flavor_pint_avgs, on='flavor').set_index(idx)\n",
    "test.pints = test_merged.pints_x.fillna(test_merged.pints_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train shape:  (4050, 2) - no. of nulls in pints 0\n",
      "validate shape:  (450, 2) - no. of nulls in pints 0\n",
      "test shape:  (500, 2) - no. of nulls in pints 0\n"
     ]
    }
   ],
   "source": [
    "print('train shape: ', train.shape, '- no. of nulls in pints', train.pints.isna().sum())\n",
    "print('validate shape: ', validate.shape, '- no. of nulls in pints', validate.pints.isna().sum())\n",
    "print('test shape: ', test.shape, '- no. of nulls in pints', test.pints.isna().sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Postscript\n",
    "\n",
    "- Make sure the group averages are calculated from the training dataset.\n",
    "- Prefer the `.replace` strategy unless you need to use multiple groups.\n",
    "- The above code for the joining example could be modified to use multiple groups instead of 1 by:\n",
    "    - When producing the group averages table, group by a list of columns instead of a single one.\n",
    "    - When merging, specify the same list of columns for the `on` keyword argument instead of a single string.\n",
    "- You could use a different aggregation other than the average (e.g. median, or mode for categorical features) by changing the way the groupby table is produced."
   ]
  }
 ],
 "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

ice_cream.csv