{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predicting NYC Taxi Fares with RAPIDS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RAPIDS is a suite of GPU accelerated data science libraries with APIs that should be familiar to users of Pandas, Dask, and Scikitlearn.\n",
    "\n",
    "This notebook focuses on showing how to use cuDF with Dask & XGBoost to scale GPU DataFrame ETL-style operations & model training out to multiple GPUs on mutliple nodes as part of Google Cloud Dataproc.\n",
    "\n",
    "Anaconda has graciously made some of the NYC Taxi dataset available in a public Google Cloud Storage bucket. We'll use our Dataproc Cluster of GPUs to process it and train a model that predicts the fare amount.\n",
    "\n",
    "For EDA we show the examples using [Holoviews](http://holoviews.org/) and [hvplot](https://hvplot.holoviz.org/). Best way to install Holoviews is to from `conda-forge` channel `conda install -c conda-forge holoviews`\n",
    "and for hvplot `pyviz` channel. `conda install -c pyviz hvplot`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import cupy\n",
    "import numba, socket, time\n",
    "import cudf\n",
    "import dask, dask_cudf\n",
    "import xgboost as xgb\n",
    "import cuspatial\n",
    "from dask.distributed import Client, wait\n",
    "from dask_cuda import LocalCUDACluster\n",
    "\n",
    "#To install Holoviews and hvplot\n",
    "#conda install -c conda-forge holoviews\n",
    "#conda install -c pyviz hvplot\n",
    "import holoviews as hv\n",
    "from holoviews import opts\n",
    "import numpy as np\n",
    "import hvplot.cudf\n",
    "import hvplot.dask\n",
    "hv.extension('bokeh')\n",
    "\n",
    "# connect to the Dask cluster created at Dataproc startup time\n",
    "# client = Client(socket.gethostname()+':8786')\n",
    "\n",
    "# if you're not using Dataproc, you can use dask_cuda to create a local CUDA cluster as well\n",
    "# as shown in the commented code below. Check https://github.com/rapidsai/dask-cuda for help\n",
    "\n",
    "cluster = LocalCUDACluster()\n",
    "client = Client(cluster)\n",
    "\n",
    "# forces workers to restart. useful to ensure GPU memory is clear\n",
    "# client.restart()\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspecting the Data\n",
    "\n",
    "Now that we have a cluster of GPU workers, we'll use [dask-cudf](https://github.com/rapidsai/dask-cudf/) to load and parse a bunch of CSV files into a distributed DataFrame. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''if you get 'ModuleNotFoundError: No module named 'gcsfs', run `!pip install gcsfs` \n",
    "'''\n",
    "base_path = '/yellow-taxi/'\n",
    "\n",
    "df_2014 = dask_cudf.read_csv(base_path+'2014/yellow_*.csv')\n",
    "df_2014.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Cleanup\n",
    "\n",
    "As usual, the data needs to be massaged a bit before we can start adding features that are useful to an ML model.\n",
    "\n",
    "For example, in the 2014 taxi CSV files, there are `pickup_datetime` and `dropoff_datetime` columns. The 2015 CSVs have `tpep_pickup_datetime` and `tpep_dropoff_datetime`, which are the same columns. One year has `rate_code`, and another `RateCodeID`.\n",
    "\n",
    "Also, some CSV files have column names with extraneous spaces in them.\n",
    "\n",
    "Worst of all, starting in the July 2016 CSVs, pickup & dropoff latitude and longitude data were replaced by location IDs, making the second half of the year useless to us.\n",
    "\n",
    "We'll do a little string manipulation, column renaming, and concatenating of DataFrames to sidestep the problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dictionary of required columns and their datatypes\n",
    "must_haves = {\n",
    "     'pickup_datetime': 'datetime64[s]',\n",
    "     'dropoff_datetime': 'datetime64[s]',\n",
    "     'passenger_count': 'int32',\n",
    "     'trip_distance': 'float32',\n",
    "     'pickup_longitude': 'float32',\n",
    "     'pickup_latitude': 'float32',\n",
    "     'rate_code': 'int32',\n",
    "     'dropoff_longitude': 'float32',\n",
    "     'dropoff_latitude': 'float32',\n",
    "     'fare_amount': 'float32'\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean(ddf, must_haves):\n",
    "    # replace the extraneous spaces in column names and lower the font type\n",
    "    tmp = {col:col.strip().lower() for col in list(ddf.columns)}\n",
    "    ddf = ddf.rename(columns=tmp)\n",
    "\n",
    "    ddf = ddf.rename(columns={\n",
    "        'tpep_pickup_datetime': 'pickup_datetime',\n",
    "        'tpep_dropoff_datetime': 'dropoff_datetime',\n",
    "        'ratecodeid': 'rate_code'\n",
    "    })\n",
    "    \n",
    "    ddf['pickup_datetime'] = ddf['pickup_datetime'].astype('datetime64[ms]')\n",
    "    ddf['dropoff_datetime'] = ddf['dropoff_datetime'].astype('datetime64[ms]')\n",
    "\n",
    "    for col in ddf.columns:\n",
    "        if col not in must_haves:\n",
    "            ddf = ddf.drop(columns=col)\n",
    "            continue\n",
    "        # if column was read as a string, recast as float\n",
    "        if ddf[col].dtype == 'object':\n",
    "            ddf[col] = ddf[col].str.fillna('-1')\n",
    "            ddf[col] = ddf[col].astype('float32')\n",
    "        else:\n",
    "            # downcast from 64bit to 32bit types\n",
    "            # Tesla T4 are faster on 32bit ops\n",
    "            if 'int' in str(ddf[col].dtype):\n",
    "                ddf[col] = ddf[col].astype('int32')\n",
    "            if 'float' in str(ddf[col].dtype):\n",
    "                ddf[col] = ddf[col].astype('float32')\n",
    "            ddf[col] = ddf[col].fillna(-1)\n",
    "    \n",
    "    return ddf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<b> NOTE: </b>We will realize that some of 2015 data has column name as `RateCodeID` and others have `RatecodeID`. When we rename the columns in the clean function, it internally doesn't pass meta while calling map_partitions(). This leads to the error of column name mismatch in the returned data. For this reason, we will call the clean function with map_partition and pass the meta to it. Here is the link to the bug created for that: https://github.com/rapidsai/cudf/issues/5413 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2014 = df_2014.map_partitions(clean, must_haves, meta=must_haves)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We still have 2015 and the first half of 2016's data to read and clean. Let's increase our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2015 = dask_cudf.read_csv(base_path+'2015/yellow_*.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2015 = df_2015.map_partitions(clean, must_haves, meta=must_haves)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Handling 2016's Mid-Year Schema Change"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In 2016, only January - June CSVs have the columns we need. If we try to read base_path+2016/yellow_*.csv, Dask will not appreciate having differing schemas in the same DataFrame.\n",
    "\n",
    "Instead, we'll need to create a list of the valid months and read them independently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "months = [str(x).rjust(2, '0') for x in range(1, 7)]\n",
    "valid_files = [base_path+'2016/yellow_tripdata_2016-'+month+'.csv' for month in months]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#read & clean 2016 data and concat all DFs\n",
    "df_2016 = dask_cudf.read_csv(valid_files).map_partitions(clean, must_haves, meta=must_haves)\n",
    "\n",
    "#concatenate multiple DataFrames into one bigger one\n",
    "taxi_df = dask.dataframe.multi.concat([df_2014, df_2015, df_2016])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_df = taxi_df.persist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploratory Data Analysis (EDA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we are checking out if there are any non-sensical records and outliers, and in such case, we need to remove them from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check out if there is any negative total trip time\n",
    "taxi_df[taxi_df.dropoff_datetime <= taxi_df.pickup_datetime].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check out if there is any abnormal data where trip distance is short, but the fare is very high.\n",
    "taxi_df[(taxi_df.trip_distance < 10) & (taxi_df.fare_amount > 300)].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check out if there is any abnormal data where trip distance is long, but the fare is very low.\n",
    "taxi_df[(taxi_df.trip_distance > 50) & (taxi_df.fare_amount < 50)].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Using only 2016-01 data for visuals.\n",
    "#taxi_df_cdf = clean(cudf.read_csv(valid_files[0]),must_haves)\n",
    "\n",
    "#Using entire 2016 data for visualization\n",
    "taxi_df_cdf = dask_cudf.read_csv(valid_files,npartitions=8).map_partitions(clean,must_haves,meta=must_haves).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plot below visualizes the histogram of trip_distance and we can see some abnormal trip_distance values for some records. Taking this and also the NYC map coordinates into consideration, we will only select records where tripdistance < 500 miles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "#Histogram using cupy and Holoviews\n",
    "# frequencies, edges = cupy.histogram(x=cupy.array(taxi_df_cdf[\"trip_distance\"]) , bins=20)\n",
    "# hist = hv.Histogram((np.array(edges.tolist()), np.array(frequencies.tolist())))\n",
    "\n",
    "#Histogram using hvplot\n",
    "hist = taxi_df_cdf.hvplot.hist(\"trip_distance\", bins=20)\n",
    "\n",
    "#Customizing the plot\n",
    "hist.opts(xlabel=\"trip distance (miles)\",ylabel=\"count\",color=\"green\",width=900, height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, the plot below visualizes the histogram of fare_amount and we can see some abnormal fare_amount values for some records. Taking this and also the NYC map coordinates into consideration, we will only select records where fare_amount < 500$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "#Histogram using cupy and Holoviews\n",
    "# frequencies, edges = cupy.histogram(x=cupy.array(taxi_df_cdf[\"fare_amount\"]) , bins=20)\n",
    "# hist = hv.Histogram((np.array(edges.tolist()), np.array(frequencies.tolist())))\n",
    "\n",
    "#Histogram using hvplot\n",
    "hist = taxi_df_cdf.hvplot.hist(\"fare_amount\", bins=20)\n",
    "\n",
    "#Customizing the plot\n",
    "hist.opts(xlabel=\"fare amount ($)\",ylabel=\"count\",color=\"green\",width=900, height=400)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Plot the number of passengers per trip. We'll remove the records where passenger_count > 5.\n",
    "# Plotting using Holoviews\n",
    "#bar = hv.Bars(taxi_df_cdf.groupby(\"passenger_count\").size().to_frame().rename(columns={0:\"count\"}))\n",
    "\n",
    "# Plotting using hvplot\n",
    "df_bar = taxi_df_cdf.groupby(\"passenger_count\").size().to_frame().rename(columns={0:\"count\"}).reset_index()\n",
    "bar = df_bar.hvplot.bar(x=\"passenger_count\",y=\"count\")\n",
    "\n",
    "#Customizing the plot\n",
    "bar.opts(color=\"green\",width=900, height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "EDA visuals and additional analysis yield the filter logic below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply a list of filter conditions to throw out records with missing or outlier values\n",
    "query_frags = [\n",
    "    'fare_amount > 1 and fare_amount < 500',\n",
    "    'passenger_count > 0 and passenger_count < 6',\n",
    "    'pickup_longitude > -75 and pickup_longitude < -73',\n",
    "    'dropoff_longitude > -75 and dropoff_longitude < -73',\n",
    "    'pickup_latitude > 40 and pickup_latitude < 42',\n",
    "    'dropoff_latitude > 40 and dropoff_latitude < 42',\n",
    "    'trip_distance > 0 and trip_distance < 500',\n",
    "    'not (trip_distance > 50 and fare_amount < 50)',\n",
    "    'not (trip_distance < 10 and fare_amount > 300)',\n",
    "    'not dropoff_datetime <= pickup_datetime'\n",
    "]\n",
    "taxi_df = taxi_df.query(' and '.join(query_frags))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset_index and drop index column\n",
    "taxi_df = taxi_df.reset_index(drop=True)\n",
    "taxi_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding Interesting Features\n",
    "\n",
    "Dask & cuDF provide standard DataFrame operations, but also let you run \"user defined functions\" on the underlying data. Here we use [dask.dataframe's map_partitions](https://docs.dask.org/en/latest/dataframe-api.html#dask.dataframe.DataFrame.map_partitions) to apply user defined python function on each DataFrame partition.\n",
    "\n",
    "We'll use a Haversine Distance calculation to find total trip distance, and extract additional useful variables from the datetime fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## add features\n",
    "\n",
    "taxi_df['hour'] = taxi_df['pickup_datetime'].dt.hour\n",
    "taxi_df['year'] = taxi_df['pickup_datetime'].dt.year\n",
    "taxi_df['month'] = taxi_df['pickup_datetime'].dt.month\n",
    "taxi_df['day'] = taxi_df['pickup_datetime'].dt.day\n",
    "taxi_df['day_of_week'] = taxi_df['pickup_datetime'].dt.weekday\n",
    "taxi_df['is_weekend'] = (taxi_df['day_of_week']>=5).astype('int32')\n",
    "\n",
    "#calculate the time difference between dropoff and pickup.\n",
    "taxi_df['diff'] = taxi_df['dropoff_datetime'].astype('int64') - taxi_df['pickup_datetime'].astype('int64')\n",
    "taxi_df['diff']=(taxi_df['diff']/1000).astype('int64')\n",
    "\n",
    "taxi_df['pickup_latitude_r'] = taxi_df['pickup_latitude']//.01*.01\n",
    "taxi_df['pickup_longitude_r'] = taxi_df['pickup_longitude']//.01*.01\n",
    "taxi_df['dropoff_latitude_r'] = taxi_df['dropoff_latitude']//.01*.01\n",
    "taxi_df['dropoff_longitude_r'] = taxi_df['dropoff_longitude']//.01*.01\n",
    "\n",
    "taxi_df = taxi_df.drop('pickup_datetime', axis=1)\n",
    "taxi_df = taxi_df.drop('dropoff_datetime', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def haversine_dist(df):\n",
    "    h_distance = cuspatial.haversine_distance(df['pickup_longitude'], df['pickup_latitude'], df['dropoff_longitude'], df['dropoff_latitude'])\n",
    "    df['h_distance']= h_distance\n",
    "    df['h_distance']= df['h_distance'].astype('float32')\n",
    "    return df\n",
    "\n",
    "taxi_df = taxi_df.map_partitions(haversine_dist)\n",
    "taxi_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_df = taxi_df.persist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pick a Training Set\n",
    "\n",
    "Let's imagine you're making a trip to New York on the 25th and want to build a model to predict what fare prices will be like the last few days of the month based on the first part of the month. We'll use a query expression to identify the `day` of the month to use to divide the data into train and test sets.\n",
    "\n",
    "The wall-time below represents how long it takes your GPU cluster to load data from the Google Cloud Storage bucket and the ETL portion of the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#since we calculated the h_distance let's drop the trip_distance column, and then do model training with XGB.\n",
    "taxi_df = taxi_df.drop('trip_distance', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the original data partition for train and test sets.\n",
    "X_train = taxi_df.query('day < 25').persist()\n",
    "\n",
    "# create a Y_train ddf with just the target variable\n",
    "Y_train = X_train[['fare_amount']].persist()\n",
    "# drop the target variable from the training ddf\n",
    "X_train = X_train[X_train.columns.difference(['fare_amount'])]\n",
    "\n",
    "# this wont return until all data is in GPU memory\n",
    "done = wait([X_train, Y_train])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the XGBoost Regression Model\n",
    "\n",
    "The wall time output below indicates how long it took your GPU cluster to train an XGBoost model over the training set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtrain = xgb.dask.DaskDMatrix(client, X_train, Y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "trained_model = xgb.dask.train(client,\n",
    "                        {\n",
    "                         'learning_rate': 0.3,\n",
    "                          'max_depth': 8,\n",
    "                          'objective': 'reg:squarederror',\n",
    "                          'subsample': 0.6,\n",
    "                          'gamma': 1,\n",
    "                          'silent': True,\n",
    "                          'verbose_eval': True,\n",
    "                          'tree_method':'gpu_hist'\n",
    "                        },\n",
    "                        dtrain,\n",
    "                        num_boost_round=100, evals=[(dtrain, 'train')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = xgb.plot_importance(trained_model['booster'], height=0.8, max_num_features=10, importance_type=\"gain\")\n",
    "ax.grid(False, axis=\"y\")\n",
    "ax.set_title('Estimated feature importance')\n",
    "ax.set_xlabel('importance')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How Good is Our Model?\n",
    "\n",
    "Now that we have a trained model, we need to test it with the 25% of records we held out.\n",
    "\n",
    "Based on the filtering conditions applied to this dataset, many of the DataFrame partitions will wind up having 0 rows. This is a problem for XGBoost which doesn't know what to do with 0 length arrays. We'll repartition the data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def drop_empty_partitions(df):\n",
    "    lengths = df.map_partitions(len).compute()\n",
    "    nonempty = [length > 0 for length in lengths]\n",
    "    return df.partitions[nonempty]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = taxi_df.query('day >= 25').persist()\n",
    "X_test = drop_empty_partitions(X_test)\n",
    "\n",
    "# Create Y_test with just the fare amount\n",
    "Y_test = X_test[['fare_amount']].persist()\n",
    "\n",
    "# Drop the fare amount from X_test\n",
    "X_test = X_test[X_test.columns.difference(['fare_amount'])]\n",
    "\n",
    "# this wont return until all data is in GPU memory\n",
    "done = wait([X_test, Y_test])\n",
    "\n",
    "# display test set size\n",
    "len(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate predictions on the test set\n",
    "'''feed X_test as a dask.dataframe'''\n",
    "\n",
    "booster = trained_model[\"booster\"] # \"Booster\" is the trained model\n",
    "history = trained_model['history'] # \"History\" is a dictionary containing evaluation results \n",
    "\n",
    "booster.set_param({'predictor': 'gpu_predictor'})\n",
    "\n",
    "prediction = xgb.dask.predict(client, booster, X_test).persist()\n",
    "\n",
    "wait(prediction)\n",
    "\n",
    "prediction.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = prediction.map_partitions(lambda part: cudf.Series(part)).reset_index(drop=True)\n",
    "actual = Y_test['fare_amount'].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: We mapped each partition of the result from `xgb.dask.predict` into `cudf.Series` to be able to substract it from `actual` data. Here is the issue asking XGBoost to solve that before returning data from `xgb.dask.predict` https://github.com/dmlc/xgboost/issues/5823#issuecomment-648526888"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actual.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate RMSE\n",
    "squared_error = ((prediction-actual)**2)\n",
    "\n",
    "# compute the actual RMSE over the full test set\n",
    "cupy.sqrt(squared_error.mean().compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Trained Model for Later Use¶"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make a model maximally useful, you need to be able to save it for later use. We'll use Google Cloud Storage to persist the trained model in a dill file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gcsfs, dill\n",
    "\n",
    "fs = gcsfs.GCSFileSystem()\n",
    "# replace with a bucket you own\n",
    "bucket = 'rapidsai-test-1/'\n",
    "\n",
    "with fs.open(bucket+'trained_model.dill', 'wb') as file:  \n",
    "    dill.dump(trained_model, file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an alternative, you can save the trained model on your system as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model to file\n",
    "booster.save_model('xgboost-model')\n",
    "print('Training evaluation history:', history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reload a Saved Model from Disk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also read the saved model back out of Google Cloud Storage and into a normal XGBoost model object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with fs.open(bucket+'trained_model.dill', 'rb') as file:  \n",
    "    model_from_disk = dill.load(file)\n",
    "\n",
    "# Generate predictions on the test set again, but this time using the reloaded model\n",
    "prediction = xgb.dask.predict(client, model_from_disk, X_test).persist()\n",
    "wait(prediction)\n",
    "\n",
    "# Verify that the predictions result in the same RMSE error\n",
    "prediction = prediction.map_partitions(lambda part: cudf.Series(part)).reset_index(drop=True)\n",
    "actual = Y_test['fare_amount'].reset_index(drop=True)\n",
    "\n",
    "squared_error = ((prediction-actual)**2)\n",
    "\n",
    "# compute the actual RMSE over the full test set\n",
    "cupy.sqrt(squared_error.mean().compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# cuML part"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Normalize data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import MinMaxScaler, StandardScaler\n",
    "scaler_x = MinMaxScaler()\n",
    "scaler_y = StandardScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler_x.fit(X_train)\n",
    "X_train = scaler_x.transform(X_train)\n",
    "X_test = scaler_x.transform(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler_y.fit(Y_train.to_numpy().reshape(-1, 1))\n",
    "Y_train = scaler_y.transform(Y_train.to_numpy().reshape(-1, 1)).ravel()\n",
    "Y_test = scaler_y.transform(Y_test.to_numpy().reshape(-1, 1)).ravel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Ridge regression model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cuml import Ridge\n",
    "from cuml.metrics.regression import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "model_r = Ridge()\n",
    "model_r.fit(X_train, Y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's look at RMSE metric of models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model_r_p.predict(X_test)\n",
    "\n",
    "mse_r = mean_squared_error(Y_test, y_pred)\n",
    "\n",
    "print(f'RMSE of Ridge: {np.sqrt(mse_r_p)}')"
   ]
  }
 ],
 "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.9.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
