Random_forest_project.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Random forest regression sample project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing the necessary libraries\n",
    "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": [
    "# reading in our data\n",
    "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": [
    "# displaying the first few columns of the population dataset\n",
    "population_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# getting a brief description of the dataset\n",
    "population_df.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It appears the population data has several missing values as follows:\n",
    "\n",
    "- 1960 to 1989 - 3 null values\n",
    "- 1992 to 1994 - 2 null values\n",
    "- 1995 to 1997 - 1 null values\n",
    "- 2012 to 2017 - 1 null values"
   ]
  },
  {
   "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": [
    "# displaying the first few rows of the metadata\n",
    "meta_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It appears that the 'Special Notes' column has 101 null values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 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": [
    "def get_total_pop_by_income(income_group_name='Low income'):\n",
    "    '''\n",
    "    This function takes the Income Group of a country as a string and fetches the year and population\n",
    "    for the entire income group returning them as an (i, 2) np.int64 array, i being the length of the\n",
    "    data. The default income group is set to 'Low Income'.\n",
    "    It also raises a ValueError if the specified value of the Income Group isn't given.\n",
    "\n",
    "    Args:\n",
    "        income_group_name (str) - The name of the Income Group.\n",
    "\n",
    "    Returns:\n",
    "        income_group_population (arr) - An array containing the years and respective population of the\n",
    "                                        given income group.\n",
    "    '''\n",
    "    merged_df = population_df.merge(meta_df, how = 'inner', on = 'Country Code')\n",
    "    \n",
    "    if income_group_name not in merged_df['Income Group'].unique():\n",
    "        raise ValueError('The given Income Group is invalid')\n",
    "    \n",
    "    income_group_df = merged_df[merged_df['Income Group'] == income_group_name]\n",
    "    total_pop = income_group_df.iloc[:,:-3].sum(axis = 0)\n",
    "    result = np.array([income_group_df.columns[:-3].astype(int), total_pop.values], dtype = np.int64).T\n",
    "    \n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[      1960,  769889923],\n",
       "       [      1961,  781225329],\n",
       "       [      1962,  791207437],\n",
       "       [      1963,  801108277],\n",
       "       [      1964,  810900987],\n",
       "       [      1965,  820309686],\n",
       "       [      1966,  829088382],\n",
       "       [      1967,  837479954],\n",
       "       [      1968,  844905494],\n",
       "       [      1969,  854059674],\n",
       "       [      1970,  862276721],\n",
       "       [      1971,  871169187],\n",
       "       [      1972,  880246152],\n",
       "       [      1973,  888486025],\n",
       "       [      1974,  897803169],\n",
       "       [      1975,  906573084],\n",
       "       [      1976,  913843314],\n",
       "       [      1977,  921330504],\n",
       "       [      1978,  928906293],\n",
       "       [      1979,  936836246],\n",
       "       [      1980,  944587066],\n",
       "       [      1981,  952368316],\n",
       "       [      1982,  959759971],\n",
       "       [      1983,  966754949],\n",
       "       [      1984,  973423742],\n",
       "       [      1985,  980143630],\n",
       "       [      1986,  987194728],\n",
       "       [      1987,  994242786],\n",
       "       [      1988, 1001421456],\n",
       "       [      1989, 1009036892],\n",
       "       [      1990, 1017092667],\n",
       "       [      1991, 1025345408],\n",
       "       [      1992, 1031949811],\n",
       "       [      1993, 1040349480],\n",
       "       [      1994, 1048121445],\n",
       "       [      1995, 1057290586],\n",
       "       [      1996, 1064630661],\n",
       "       [      1997, 1071969568],\n",
       "       [      1998, 1078927765],\n",
       "       [      1999, 1085992668],\n",
       "       [      2000, 1092825678],\n",
       "       [      2001, 1100293969],\n",
       "       [      2002, 1107836355],\n",
       "       [      2003, 1115390519],\n",
       "       [      2004, 1123325037],\n",
       "       [      2005, 1131426281],\n",
       "       [      2006, 1140084827],\n",
       "       [      2007, 1149238990],\n",
       "       [      2008, 1158965286],\n",
       "       [      2009, 1167712409],\n",
       "       [      2010, 1175649232],\n",
       "       [      2011, 1181451343],\n",
       "       [      2012, 1188796100],\n",
       "       [      2013, 1196212921],\n",
       "       [      2014, 1203819897],\n",
       "       [      2015, 1211252041],\n",
       "       [      2016, 1218629612],\n",
       "       [      2017, 1225514228]], dtype=int64)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = get_total_pop_by_income('High income')\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 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). \n",
    "\n",
    "Using this knowledge, we 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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sklearn_kfold_split(data,K):\n",
    "    '''\n",
    "    This function takes an array and splits it into training and testind data, returning them\n",
    "    as K number of tuples, K being an integer input for the Kfold class from sklearn.\n",
    "\n",
    "    Args:\n",
    "        data (arr) - The data to be splitted\n",
    "        K (int) - The number of times the data is to be split\n",
    "\n",
    "    Returns:\n",
    "        data_indices (list) - Training and testing indices corresponding to the data points\n",
    "                            that belong to a particular split\n",
    "    '''\n",
    "    kf = KFold(n_splits=K, shuffle=False)\n",
    "    data_indices = ([])\n",
    "    \n",
    "    for train_indices, test_indices in kf.split(data):\n",
    "        data_indices.append((train_indices.tolist(), test_indices.tolist()))\n",
    "    \n",
    "    return data_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "list"
      ]
     },
     "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": [
    "### Step 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, we 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def best_k_model(data,data_indices):\n",
    "    '''\n",
    "    This function finds the best model fit on the data given.\n",
    "\n",
    "    Args:\n",
    "        data (arr) - An array containing the data for both traning and testing\n",
    "        data_indices (list) - Training and testing indices corresponding to the data points\n",
    "                            that belong to a particular split\n",
    "\n",
    "    Returns:\n",
    "        best_model (model) - The model with the least MSE\n",
    "    '''\n",
    "    best_model = None\n",
    "    best_mse = float('inf')  # Initialize with a large value\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",
    "        # Train RandomForestRegressor model\n",
    "        rf = RandomForestRegressor(random_state=42)\n",
    "        rf.fit(X_train.reshape(-1, 1), y_train)\n",
    "        \n",
    "        # Evaluate model on test set\n",
    "        y_pred = rf.predict(X_test.reshape(-1, 1))\n",
    "        mse = mean_squared_error(y_test, y_pred)\n",
    "        \n",
    "        # Update best model if current model has lower mse\n",
    "        if mse < best_mse:\n",
    "            best_mse = mse\n",
    "            best_model = rf\n",
    "    \n",
    "    return best_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([7.76089629e+08])"
      ]
     },
     "execution_count": 14,
     "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]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}