11.08.the_random_forest_student_version.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"center\" style=\" font-size: 80%; text-align: center; margin: 0 auto\">\n",
    "<img src=\"https://raw.githubusercontent.com/Explore-AI/Pictures/master/Python-Notebook-Banners/Code_challenge.png\"  style=\"display: block; margin-left: auto; margin-right: auto;\";/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code challenge: Random forest regression\n",
    "© ExploreAI Academy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this code challenge, we'll test our knowledge of how to create an ensemble model known as a random forest. We will train this new model using the world population data. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "⚠️ **Note that this code challenge is graded and will contribute to your overall marks for this module. Submit this notebook for grading. Note that the names of the functions are different in this notebook. Transfer the code in your notebook to this submission notebook**\n",
    "\n",
    "### Instructions\n",
    "\n",
    "- **Do not add or remove cells in this notebook. Do not edit or remove the `### START FUNCTION` or `### END FUNCTION` comments. Do not add any code outside of the functions you are required to edit. Doing any of this will lead to a mark of 0%!**\n",
    "\n",
    "- Answer the questions according to the specifications provided.\n",
    "\n",
    "- Use the given cell in each question to see if your function matches the expected outputs.\n",
    "\n",
    "- Do not hard-code answers to the questions.\n",
    "\n",
    "- The use of StackOverflow, Google, and other online tools is permitted. However, copying a fellow student's code is not permissible and is considered a breach of the Honour code. Doing this will result in a mark of 0%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from numpy import array\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "population_df = pd.read_csv('https://raw.githubusercontent.com/Explore-AI/Public-Data/master/AnalyseProject/world_population.csv', index_col='Country Code')\n",
    "meta_df = pd.read_csv('https://raw.githubusercontent.com/Explore-AI/Public-Data/master/AnalyseProject/metadata.csv', index_col='Country Code')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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>1960</th>\n",
       "      <th>1961</th>\n",
       "      <th>1962</th>\n",
       "      <th>1963</th>\n",
       "      <th>1964</th>\n",
       "      <th>1965</th>\n",
       "      <th>1966</th>\n",
       "      <th>1967</th>\n",
       "      <th>1968</th>\n",
       "      <th>1969</th>\n",
       "      <th>...</th>\n",
       "      <th>2008</th>\n",
       "      <th>2009</th>\n",
       "      <th>2010</th>\n",
       "      <th>2011</th>\n",
       "      <th>2012</th>\n",
       "      <th>2013</th>\n",
       "      <th>2014</th>\n",
       "      <th>2015</th>\n",
       "      <th>2016</th>\n",
       "      <th>2017</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Country Code</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ABW</th>\n",
       "      <td>54211.0</td>\n",
       "      <td>55438.0</td>\n",
       "      <td>56225.0</td>\n",
       "      <td>56695.0</td>\n",
       "      <td>57032.0</td>\n",
       "      <td>57360.0</td>\n",
       "      <td>57715.0</td>\n",
       "      <td>58055.0</td>\n",
       "      <td>58386.0</td>\n",
       "      <td>58726.0</td>\n",
       "      <td>...</td>\n",
       "      <td>101353.0</td>\n",
       "      <td>101453.0</td>\n",
       "      <td>101669.0</td>\n",
       "      <td>102053.0</td>\n",
       "      <td>102577.0</td>\n",
       "      <td>103187.0</td>\n",
       "      <td>103795.0</td>\n",
       "      <td>104341.0</td>\n",
       "      <td>104822.0</td>\n",
       "      <td>105264.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AFG</th>\n",
       "      <td>8996351.0</td>\n",
       "      <td>9166764.0</td>\n",
       "      <td>9345868.0</td>\n",
       "      <td>9533954.0</td>\n",
       "      <td>9731361.0</td>\n",
       "      <td>9938414.0</td>\n",
       "      <td>10152331.0</td>\n",
       "      <td>10372630.0</td>\n",
       "      <td>10604346.0</td>\n",
       "      <td>10854428.0</td>\n",
       "      <td>...</td>\n",
       "      <td>27294031.0</td>\n",
       "      <td>28004331.0</td>\n",
       "      <td>28803167.0</td>\n",
       "      <td>29708599.0</td>\n",
       "      <td>30696958.0</td>\n",
       "      <td>31731688.0</td>\n",
       "      <td>32758020.0</td>\n",
       "      <td>33736494.0</td>\n",
       "      <td>34656032.0</td>\n",
       "      <td>35530081.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AGO</th>\n",
       "      <td>5643182.0</td>\n",
       "      <td>5753024.0</td>\n",
       "      <td>5866061.0</td>\n",
       "      <td>5980417.0</td>\n",
       "      <td>6093321.0</td>\n",
       "      <td>6203299.0</td>\n",
       "      <td>6309770.0</td>\n",
       "      <td>6414995.0</td>\n",
       "      <td>6523791.0</td>\n",
       "      <td>6642632.0</td>\n",
       "      <td>...</td>\n",
       "      <td>21759420.0</td>\n",
       "      <td>22549547.0</td>\n",
       "      <td>23369131.0</td>\n",
       "      <td>24218565.0</td>\n",
       "      <td>25096150.0</td>\n",
       "      <td>25998340.0</td>\n",
       "      <td>26920466.0</td>\n",
       "      <td>27859305.0</td>\n",
       "      <td>28813463.0</td>\n",
       "      <td>29784193.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ALB</th>\n",
       "      <td>1608800.0</td>\n",
       "      <td>1659800.0</td>\n",
       "      <td>1711319.0</td>\n",
       "      <td>1762621.0</td>\n",
       "      <td>1814135.0</td>\n",
       "      <td>1864791.0</td>\n",
       "      <td>1914573.0</td>\n",
       "      <td>1965598.0</td>\n",
       "      <td>2022272.0</td>\n",
       "      <td>2081695.0</td>\n",
       "      <td>...</td>\n",
       "      <td>2947314.0</td>\n",
       "      <td>2927519.0</td>\n",
       "      <td>2913021.0</td>\n",
       "      <td>2905195.0</td>\n",
       "      <td>2900401.0</td>\n",
       "      <td>2895092.0</td>\n",
       "      <td>2889104.0</td>\n",
       "      <td>2880703.0</td>\n",
       "      <td>2876101.0</td>\n",
       "      <td>2873457.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AND</th>\n",
       "      <td>13411.0</td>\n",
       "      <td>14375.0</td>\n",
       "      <td>15370.0</td>\n",
       "      <td>16412.0</td>\n",
       "      <td>17469.0</td>\n",
       "      <td>18549.0</td>\n",
       "      <td>19647.0</td>\n",
       "      <td>20758.0</td>\n",
       "      <td>21890.0</td>\n",
       "      <td>23058.0</td>\n",
       "      <td>...</td>\n",
       "      <td>83861.0</td>\n",
       "      <td>84462.0</td>\n",
       "      <td>84449.0</td>\n",
       "      <td>83751.0</td>\n",
       "      <td>82431.0</td>\n",
       "      <td>80788.0</td>\n",
       "      <td>79223.0</td>\n",
       "      <td>78014.0</td>\n",
       "      <td>77281.0</td>\n",
       "      <td>76965.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 58 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                   1960       1961       1962       1963       1964  \\\n",
       "Country Code                                                          \n",
       "ABW             54211.0    55438.0    56225.0    56695.0    57032.0   \n",
       "AFG           8996351.0  9166764.0  9345868.0  9533954.0  9731361.0   \n",
       "AGO           5643182.0  5753024.0  5866061.0  5980417.0  6093321.0   \n",
       "ALB           1608800.0  1659800.0  1711319.0  1762621.0  1814135.0   \n",
       "AND             13411.0    14375.0    15370.0    16412.0    17469.0   \n",
       "\n",
       "                   1965        1966        1967        1968        1969  ...  \\\n",
       "Country Code                                                             ...   \n",
       "ABW             57360.0     57715.0     58055.0     58386.0     58726.0  ...   \n",
       "AFG           9938414.0  10152331.0  10372630.0  10604346.0  10854428.0  ...   \n",
       "AGO           6203299.0   6309770.0   6414995.0   6523791.0   6642632.0  ...   \n",
       "ALB           1864791.0   1914573.0   1965598.0   2022272.0   2081695.0  ...   \n",
       "AND             18549.0     19647.0     20758.0     21890.0     23058.0  ...   \n",
       "\n",
       "                    2008        2009        2010        2011        2012  \\\n",
       "Country Code                                                               \n",
       "ABW             101353.0    101453.0    101669.0    102053.0    102577.0   \n",
       "AFG           27294031.0  28004331.0  28803167.0  29708599.0  30696958.0   \n",
       "AGO           21759420.0  22549547.0  23369131.0  24218565.0  25096150.0   \n",
       "ALB            2947314.0   2927519.0   2913021.0   2905195.0   2900401.0   \n",
       "AND              83861.0     84462.0     84449.0     83751.0     82431.0   \n",
       "\n",
       "                    2013        2014        2015        2016        2017  \n",
       "Country Code                                                              \n",
       "ABW             103187.0    103795.0    104341.0    104822.0    105264.0  \n",
       "AFG           31731688.0  32758020.0  33736494.0  34656032.0  35530081.0  \n",
       "AGO           25998340.0  26920466.0  27859305.0  28813463.0  29784193.0  \n",
       "ALB            2895092.0   2889104.0   2880703.0   2876101.0   2873457.0  \n",
       "AND              80788.0     79223.0     78014.0     77281.0     76965.0  \n",
       "\n",
       "[5 rows x 58 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "population_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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>Region</th>\n",
       "      <th>Income Group</th>\n",
       "      <th>Special Notes</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Country Code</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ABW</th>\n",
       "      <td>Latin America &amp; Caribbean</td>\n",
       "      <td>High income</td>\n",
       "      <td>Mining is included in agriculture\\r\\r\\r\\nElect...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AFG</th>\n",
       "      <td>South Asia</td>\n",
       "      <td>Low income</td>\n",
       "      <td>Fiscal year end: March 20; reporting period fo...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AGO</th>\n",
       "      <td>Sub-Saharan Africa</td>\n",
       "      <td>Lower middle income</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ALB</th>\n",
       "      <td>Europe &amp; Central Asia</td>\n",
       "      <td>Upper middle income</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>AND</th>\n",
       "      <td>Europe &amp; Central Asia</td>\n",
       "      <td>High income</td>\n",
       "      <td>WB-3 code changed from ADO to AND to align wit...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                 Region         Income Group  \\\n",
       "Country Code                                                   \n",
       "ABW           Latin America & Caribbean          High income   \n",
       "AFG                          South Asia           Low income   \n",
       "AGO                  Sub-Saharan Africa  Lower middle income   \n",
       "ALB               Europe & Central Asia  Upper middle income   \n",
       "AND               Europe & Central Asia          High income   \n",
       "\n",
       "                                                  Special Notes  \n",
       "Country Code                                                     \n",
       "ABW           Mining is included in agriculture\\r\\r\\r\\nElect...  \n",
       "AFG           Fiscal year end: March 20; reporting period fo...  \n",
       "AGO                                                         NaN  \n",
       "ALB                                                         NaN  \n",
       "AND           WB-3 code changed from ADO to AND to align wit...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The world population data spans from 1960 to 2017. We'd like to build a predictive model that can give us the best guess at what the world population in a given year was. However, we want to compute this estimate for only _countries within a given income group_. \n",
    "\n",
    "First, however, we need to organise our data such that the sklearn's `RandomForestRegressor` class can train on our data. To do this, we will write a function that takes as input an income group and returns a 2-d numpy array that contains the year and the measured population.\n",
    "\n",
    "_**Function Specifications:**_\n",
    "* Should take a `str` argument, called `income_group_name` as input and return a numpy `array` type as output.\n",
    "* Set the default argument of `income_group_name` to equal `'Low income'`.\n",
    "* If the specified value of `income_group_name` does not exist, the function must raise a `ValueError`.\n",
    "* The array should only have two columns containing the year and the population, in other words, it should have a shape `(?, 2)` where `?` is the length of the data.\n",
    "* The values within the array should be of type `np.int64`. \n",
    "\n",
    "_**Further Reading:**_\n",
    "\n",
    "Data types are associated with memory allocation. As such, your choice of data type affects the precision of computations in your program. For example, the `np.int` data type in numpy can only store values between -2147483648 to 2147483647 and assigning values outside this range for variables of this data type may cause run-time errors. To avoid this, we can use data types with larger memory capacity e.g. `np.int64`.\n",
    "\n",
    "https://docs.scipy.org/doc/numpy/user/basics.types.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "### START FUNCTION\n",
    "def get_total_pop_by_income(income_group_name='Low income'):\n",
    "  if meta_df[meta_df['Income Group'] == income_group_name].shape[0] == 0:\n",
    "    raise ValueError\n",
    "  else:\n",
    "    selected_countries = meta_df[meta_df['Income Group'] == income_group_name].index\n",
    "\n",
    "  array = []\n",
    "  sel_df = population_df.loc[selected_countries].sum().astype('Int64')\n",
    "  for i in range(len(sel_df)) :\n",
    "    array.append([int(sel_df.index[i]), sel_df.values[i]])\n",
    "  array2 = np.array(array)\n",
    "  return array2\n",
    "\n",
    "### END FUNCTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_total_pop_by_income('High income')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_**Expected Outputs:**_\n",
    "```python\n",
    "get_total_pop_by_income('High income')\n",
    "```\n",
    "> ```\n",
    "array([[      1960,  769889923],\n",
    "       [      1961,  781225329],\n",
    "       [      1962,  791207437],\n",
    "       [      1963,  801108277],\n",
    "       ...\n",
    "       [      2015, 1211252041],\n",
    "       [      2016, 1218629612],\n",
    "       [      2017, 1225514228]])\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 2\n",
    "\n",
    "Now that we have have our data, we need to split this into a set of variables we will be training on, and the set of variables that we will make our predictions on.\n",
    "\n",
    "Sklearn has a bunch of built-in functionality for creating training and testing sets. Our task is to implement a k-fold cross validation split of the data using sklearn's `KFold` [class](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.KFold.html) (which has already been imported into this notebook for your convenience). \n",
    "\n",
    "Using this knowledge, write a function which uses sklearn's `KFold` [class](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.KFold.html) internally, and that will take as input a 2-d numpy array and an integer `K` corresponding to the number of splits. This function will then return a list of tuples of length `K`. Each tuple in this list should consist of a `train_indices` list and a `test_indices` list containing the training/testing data point indices for that particular $K^{th}$ split.\n",
    "\n",
    "_**Function Specifications:**_\n",
    "* Should take a 2-d numpy `array` and an integer `K` as input.\n",
    "* Should use sklearn's `KFold` [class](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.KFold.html).\n",
    "* Should return a list of `K` `tuples` containing a list of training and testing indices corresponding to the data points that belong to a particular split. For example, given an array called `data` and an integer `K`, the function should return: \n",
    ">```\n",
    "data_indices = [(list_of_train_indices_for_split_1, list_of_test_indices_for_split_1)\n",
    "                  (list_of_train_indices_for_split_2, list_of_test_indices_for_split_2)\n",
    "                  (list_of_train_indices_for_split_3, list_of_test_indices_for_split_3)\n",
    "                                                   ...\n",
    "                                                   ...\n",
    "                  (list_of_train_indices_for_split_K, list_of_test_indices_for_split_K)]\n",
    "```\n",
    "\n",
    "* The `shuffle` argument in the KFold object should be set to `False`.\n",
    "\n",
    "**_Hint_**: To see an example of how to use the `KFold` object enter `help(KFold)` in a new notebook cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "### START FUNCTION\n",
    "def sklearn_kfold_split(data,K):\n",
    "  folds = []\n",
    "  kfold = KFold(n_splits=K, shuffle=False)\n",
    "  for train_index, test_index in kfold.split(data):\n",
    "    folds.append((train_index.tolist(), test_index.tolist()))\n",
    "  return folds\n",
    "\n",
    "### END FUNCTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[([15,\n",
       "   16,\n",
       "   17,\n",
       "   18,\n",
       "   19,\n",
       "   20,\n",
       "   21,\n",
       "   22,\n",
       "   23,\n",
       "   24,\n",
       "   25,\n",
       "   26,\n",
       "   27,\n",
       "   28,\n",
       "   29,\n",
       "   30,\n",
       "   31,\n",
       "   32,\n",
       "   33,\n",
       "   34,\n",
       "   35,\n",
       "   36,\n",
       "   37,\n",
       "   38,\n",
       "   39,\n",
       "   40,\n",
       "   41,\n",
       "   42,\n",
       "   43,\n",
       "   44,\n",
       "   45,\n",
       "   46,\n",
       "   47,\n",
       "   48,\n",
       "   49,\n",
       "   50,\n",
       "   51,\n",
       "   52,\n",
       "   53,\n",
       "   54,\n",
       "   55,\n",
       "   56,\n",
       "   57],\n",
       "  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]),\n",
       " ([0,\n",
       "   1,\n",
       "   2,\n",
       "   3,\n",
       "   4,\n",
       "   5,\n",
       "   6,\n",
       "   7,\n",
       "   8,\n",
       "   9,\n",
       "   10,\n",
       "   11,\n",
       "   12,\n",
       "   13,\n",
       "   14,\n",
       "   30,\n",
       "   31,\n",
       "   32,\n",
       "   33,\n",
       "   34,\n",
       "   35,\n",
       "   36,\n",
       "   37,\n",
       "   38,\n",
       "   39,\n",
       "   40,\n",
       "   41,\n",
       "   42,\n",
       "   43,\n",
       "   44,\n",
       "   45,\n",
       "   46,\n",
       "   47,\n",
       "   48,\n",
       "   49,\n",
       "   50,\n",
       "   51,\n",
       "   52,\n",
       "   53,\n",
       "   54,\n",
       "   55,\n",
       "   56,\n",
       "   57],\n",
       "  [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]),\n",
       " ([0,\n",
       "   1,\n",
       "   2,\n",
       "   3,\n",
       "   4,\n",
       "   5,\n",
       "   6,\n",
       "   7,\n",
       "   8,\n",
       "   9,\n",
       "   10,\n",
       "   11,\n",
       "   12,\n",
       "   13,\n",
       "   14,\n",
       "   15,\n",
       "   16,\n",
       "   17,\n",
       "   18,\n",
       "   19,\n",
       "   20,\n",
       "   21,\n",
       "   22,\n",
       "   23,\n",
       "   24,\n",
       "   25,\n",
       "   26,\n",
       "   27,\n",
       "   28,\n",
       "   29,\n",
       "   44,\n",
       "   45,\n",
       "   46,\n",
       "   47,\n",
       "   48,\n",
       "   49,\n",
       "   50,\n",
       "   51,\n",
       "   52,\n",
       "   53,\n",
       "   54,\n",
       "   55,\n",
       "   56,\n",
       "   57],\n",
       "  [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43]),\n",
       " ([0,\n",
       "   1,\n",
       "   2,\n",
       "   3,\n",
       "   4,\n",
       "   5,\n",
       "   6,\n",
       "   7,\n",
       "   8,\n",
       "   9,\n",
       "   10,\n",
       "   11,\n",
       "   12,\n",
       "   13,\n",
       "   14,\n",
       "   15,\n",
       "   16,\n",
       "   17,\n",
       "   18,\n",
       "   19,\n",
       "   20,\n",
       "   21,\n",
       "   22,\n",
       "   23,\n",
       "   24,\n",
       "   25,\n",
       "   26,\n",
       "   27,\n",
       "   28,\n",
       "   29,\n",
       "   30,\n",
       "   31,\n",
       "   32,\n",
       "   33,\n",
       "   34,\n",
       "   35,\n",
       "   36,\n",
       "   37,\n",
       "   38,\n",
       "   39,\n",
       "   40,\n",
       "   41,\n",
       "   42,\n",
       "   43],\n",
       "  [44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57])]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = get_total_pop_by_income('High income');\n",
    "sklearn_kfold_split(data,4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_**Expected Outputs:**_\n",
    "```python\n",
    "data = get_total_pop_by_income('High income')\n",
    "sklearn_kfold_split(data,4)\n",
    "```\n",
    "> ```\n",
    "[(array([15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,\n",
    "         32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,\n",
    "         49, 50, 51, 52, 53, 54, 55, 56, 57]),\n",
    "  array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14])),\n",
    " (array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 30, 31,\n",
    "         32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,\n",
    "         49, 50, 51, 52, 53, 54, 55, 56, 57]),\n",
    "  array([15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29])),\n",
    " (array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,\n",
    "         17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 44, 45, 46, 47,\n",
    "         48, 49, 50, 51, 52, 53, 54, 55, 56, 57]),\n",
    "  array([30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43])),\n",
    " (array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,\n",
    "         17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,\n",
    "         34, 35, 36, 37, 38, 39, 40, 41, 42, 43]),\n",
    "  array([44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57]))]\n",
    " ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question 3\n",
    "\n",
    "Now that we have formatted our data, we can fit a model using sklearn's `RandomForestRegressor` class. We'll write a function that will take as input the data indices (consisting of the train and test indices for each split) that we created in the last question, train a different `RandomForestRegressor` on each split and return the model that obtains the best testing set performance across all K splits.\n",
    "\n",
    "**Important Note:** Due to the random initialisation process used within sklearn's `RandomForestRegressor` class, you will need to fix the value of the `random_state` argument in order to get repeatable and predictable results.\n",
    "\n",
    "_**Function Specifications:**_\n",
    "* Should take a 2-d numpy array (i.e. the data) and `data_indices` (a list of `(train_indices,test_indices)` tuples) as input.\n",
    "* For each `(train_indices,test_indices)` tuple in `data_indices` the function should:\n",
    "    * Train a new `RandomForestRegressor` model on the portion of data indexed by `train_indices`\n",
    "    * Evaluate the trained `RandomForestRegressor` model on the portion of data indexed by `test_indices` using the **mean squared error** (which has also been imported for your convenience).\n",
    "* After training and evaluating the `RandomForestRegressor` models, the function should return the `RandomForestRegressor` model that obtained the highest testing set `mean_square_error` over its allocated data split across all trained models. \n",
    "* The trained `RandomForestRegressor` models should be trained with `random_state` equal `42`, all other parameters should be left as default.\n",
    "\n",
    "**_Hint_**: for each tuple in the `data_indices` list, you can obtain `X_train`,`X_test`, `y_train`, `y_test` as follows:  \n",
    ">```\n",
    "    X_train, y_train = data[train_indices,0],data[train_indices,1]\n",
    "    X_test, y_test = data[test_indices,0],data[test_indices,1]\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "### START FUNCTION\n",
    "def best_k_model(data,data_indices):\n",
    "    best_model = None\n",
    "    best_mse = float(0)\n",
    "\n",
    "    for train_indices, test_indices in data_indices:\n",
    "        x_train, y_train = data[train_indices, 0], data[train_indices, 1]\n",
    "        x_test, y_test = data[test_indices, 0], data[test_indices, 1]\n",
    "        \n",
    "        model = RandomForestRegressor(random_state=42)\n",
    "        model.fit(x_train.reshape(-1, 1), y_train)\n",
    "\n",
    "        y_pred = model.predict(x_test.reshape(-1,1))\n",
    "        mse = mean_squared_error(y_test, y_pred)\n",
    "\n",
    "        if mse > best_mse:\n",
    "            best_mse = mse\n",
    "            best_model = model\n",
    "\n",
    "    return best_model\n",
    "\n",
    "### END FUNCTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([8.85170916e+08])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = get_total_pop_by_income('High income')\n",
    "data_indices = sklearn_kfold_split(data,5)\n",
    "\n",
    "best_model = best_k_model(data,data_indices)\n",
    "best_model.predict([[1960]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_**Expected Outputs:**_\n",
    "```python\n",
    "best_model.predict([[1960]]) == array([8.85170916e+08])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  \n",
    "\n",
    "<div align=\"center\" style=\" font-size: 80%; text-align: center; margin: 0 auto\">\n",
    "<img src=\"https://raw.githubusercontent.com/Explore-AI/Pictures/master/ExploreAI_logos/EAI_Blue_Dark.png\"  style=\"width:200px\";/>\n",
    "</div>"
   ]
  }
 ],
 "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}