{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&ensp;\n",
    "[Home Page](../START_HERE.ipynb)\n",
    "\n",
    "[Previous Notebook](01-Intro_to_cuDF.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](01-Intro_to_cuDF.ipynb)\n",
    "[2]\n",
    "[3](03-Cudf_Exercise.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[Next Notebook](03-Cudf_Exercise.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User Defined Functions with cuDF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, the built-in methods of cudf.DataFrame don't do exactly what we want. We need to write a custom function (also known as a user defined function) to apply over the DataFrame.\n",
    "\n",
    "cuDF’s DataFrame class has two primary methods that let users run custom Python functions on GPUs: `apply_rows` and `apply_chunks`. In this tutorial, we’ll walk through how to use `apply_rows` and `apply_chunks` to create your own UDFs and show how you can implement a GPU-accelerated windowing function. At the end, we'll also walk through a more advanced example of applying a user defined function on a grouped DataFrame (using `apply_grouped`). The tutorial is split into different modules as shown in the list below, based on their individual content and also contain embedded exercises for your practice.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Here is the list of contents in the lab:\n",
    "- <a href='#rows'>apply_rows Method</a><br> Learn how to write a kernel function to process each of the Dataframe rows independently in parallel.\n",
    "- <a href='#chunks'>apply_chunks Method</a><br> Learn how to write a kernel function with control over how the data is chunked on the GPU.\n",
    "- <a href='#groups'>apply_grouped Method</a><br> Learn how to write a kernel function to deal with multiple time series that we want to process logically separately."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='rows'></a><br>\n",
    "\n",
    "## `apply_rows`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`apply_rows` processes each of the DataFrame rows independently in parallel. Under the hood, the `apply_rows` method will optimally divide the long columns into chunks, and assign chunks into different GPU blocks for parallel computation. \n",
    "\n",
    "In order to use `apply_rows`, we need to write a kernel function. A kernel function is a function that will be executed on each row of the DataFrame set the output value for each row. **The execution order of rows is arbitrary, so each execution of the function MUST be independent of other execution.**\n",
    "\n",
    "How does this work? User defined functions with cuDF rely on CUDA under the hood. Exploring CUDA and GPU architecture in-depth is out of scope for this tutorial. But, at a very high level, in cuDF's user defined functions:\n",
    "\n",
    "\n",
    "- Compute is spread across multiple \"blocks\", which have access to both global memory but also their own in-block memory \n",
    "- Within each block, many \"threads\" operate independently and can quickly access data in their block-specific shared memory \n",
    "\n",
    "\n",
    "As a result, the loop in the example function below resembles serial code, but executes in parallel in multiple threads on the GPU. When `kernel` is invoked, the function arguments corresponding to the input/output are strided so as to improve GPU parallelism. The kernel function is compiled to the GPU using `numba.cuda`, so the kernel function must only use Python features/functions that are [supported](https://numba.pydata.org/numba-doc/dev/cuda/cudapysupported.html) by Numba for CUDA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A kernel function takes the form\n",
    "\n",
    "```python\n",
    "def kernel(in1, in2, in3, ..., out1, out2, ..., kwarg1, kwarg2, ...):\n",
    "    for i, (x, y, z, ...) in enumerate(zip(in1, in2, in3, ...)):\n",
    "        out1[i] = ...\n",
    "        out2[i] = ...\n",
    "```\n",
    "\n",
    "`in1, in2, in3, ...` are the input columns. `out1, out2, ...` are the output columns. The kernel function should not return a result. Instead, output columns are passed as arguments and the result is written to them. Each thread writes a result to a specific index in the output column, which is why write `out1[i] = ...`.\n",
    "\n",
    "Additional keyword arguments can be passed (`kwarg1, kwarg2, ...`). Inside the kernel function, [standard numba.cuda attributes](https://numba.pydata.org/numba-doc/dev/cuda/kernels.html#thread-positioning) like `numba.cuda.threadIdx` can be used to access things like the thread or block indices. We'll explain a little bit more about this below.\n",
    "\n",
    "To execute this function on our DataFrame, we use `apply_rows`. `apply_rows` is called like:\n",
    "\n",
    "```python\n",
    "df = df.apply_rows(kernel\n",
    "                   incols=['in1', 'in2', 'in3', ...],\n",
    "                   outcols={'out1': np.float64, 'out2': np.int8, ...},\n",
    "                   kwargs={'kwarg1': val1, 'kwarg2': val2, ...})\n",
    "```\n",
    "\n",
    "`incols` is a list of the arguments for our `kernel` function representing the columns in the DataFrame. As a result, `in1`, `in2`, etc. must match the names of columns in the DataFrame that we intend to use.\n",
    "\n",
    "`outcols` is a dictionary mapping the output column names to their dtype. If we intend to generate two output columns, we need to `outcols` needs to contain two keys (`out1` and `out2`).\n",
    "\n",
    "`kwargs` is a dictionary mapping the keyword argument parameters to their values. If our `kernel` function needs additional arguments contained in our DataFrame, we can pass them in here. `kwargs` can be an empty dictionary if there are no keyword arguments.\n",
    "\n",
    "After calling `apply_rows` as above, `df` would have extra columns `out1`, `out2`, ... with the output results. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Haversine distance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example below, we create a DataFrame representing pairs of latitude and longitude points. We use `apply_rows` to calculate the [Haversine distance](https://en.wikipedia.org/wiki/Haversine_formula) between two points in the input arrays.\n",
    "\n",
    "$$\n",
    "d = 2r \\arcsin\\left(\\sqrt{\\sin^2\\left(\\frac{\\varphi_2 - \\varphi_1}{2}\\right) + \\cos(\\varphi_1) \\cos(\\varphi_2)\\sin^2\\left(\\frac{\\lambda_2 - \\lambda_1}{2}\\right)}\\right)\n",
    "$$\n",
    "\n",
    "where $\\varphi_1,\\varphi_2$ are the latitudes and $\\lambda_1,\\lambda_2$ are the longitudes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import cos, sin, asin, sqrt, pi, atan2\n",
    "\n",
    "import time\n",
    "import cudf\n",
    "import numpy as np\n",
    "from numba import cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(12)\n",
    "data_length = 1000\n",
    "\n",
    "df = cudf.DataFrame()\n",
    "df['lat1'] = np.random.normal(10, 1, data_length)\n",
    "df['lon1'] = np.random.normal(10, 1, data_length)\n",
    "df['lat2'] = np.random.normal(10, 1, data_length)\n",
    "df['lon2'] = np.random.normal(10, 1, data_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def haversine_distance_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    \"\"\"Haversine distance formula taken from Michael Dunn's StackOverflow post:\n",
    "    https://stackoverflow.com/questions/4913349/haversine-formula-in-python-bearing-and-distance-between-two-gps-points\n",
    "    \"\"\"\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread id:', cuda.threadIdx.x, 'block id:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "        \n",
    "        dlon = y_2 - y_1\n",
    "        dlat = x_2 - x_1\n",
    "        a = sin(dlat/2)**2 + cos(x_1) * cos(x_2) * sin(dlon/2)**2\n",
    "        \n",
    "        c = 2 * asin(sqrt(a)) \n",
    "        r = 6371 # Radius of earth in kilometers\n",
    "        \n",
    "        out[i] = c * r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "df = df.apply_rows(haversine_distance_kernel,\n",
    "                   incols=['lat1', 'lon1', 'lat2', 'lon2'],\n",
    "                   outcols=dict(out=np.float64),\n",
    "                   kwargs=dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we had a `print` statement in our kernel but didn't see any printed output. Print statements in kernels will only appear in terminal output; Jupyter Notebooks won't display them. We included this for this tutorial, and have copied some sample output below:\n",
    "\n",
    "```\n",
    "...\n",
    "thread id: 0 block id: 4 array size: 1 block threads: 64\n",
    "thread id: 1 block id: 4 array size: 1 block threads: 64\n",
    "thread id: 2 block id: 4 array size: 1 block threads: 64\n",
    "...\n",
    "thread id: 61 block id: 4 array size: 1 block threads: 64\n",
    "thread id: 62 block id: 4 array size: 1 block threads: 64\n",
    "thread id: 63 block id: 4 array size: 1 block threads: 64\n",
    "...\n",
    "thread id: 29 block id: 0 array size: 2 block threads: 64\n",
    "thread id: 30 block id: 0 array size: 2 block threads: 64\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, the printed output from applying our `haversine_distance_kernel` function shows some informative information. If you were to look at the entire printed output, you'd notice a few things:\n",
    "\n",
    "- The processing was spread across 15 CUDA blocks\n",
    "- Within each block, 64 separate threads were used for computation.\n",
    "- In this case, most threads in a block handled one element from the input array, but some threads have to deal with two elements, because there are 1000 rows and 960 threads (15 blocks * 64 threads per block)\n",
    "\n",
    "`apply_rows` handled all of this for us!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='exercise1'></a>\n",
    "\n",
    "\n",
    "## **Exercise 1**\n",
    "\n",
    "Modify the above example to pass in the radius of the earth `r` as a keyword argument to the kernel.\n",
    "\n",
    "<details><summary><b>Solution</b></summary>\n",
    "   <pre>\n",
    "def haversine_distance_kernel(lat1, lon1, lat2, lon2, out, r):\n",
    "    \"\"\"Haversine distance formula taken from Michael Dunn's StackOverflow post:\n",
    "    https://stackoverflow.com/questions/4913349/haversine-formula-in-python-bearing-and-distance-between-two-gps-points\n",
    "    \"\"\"\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread_id:', cuda.threadIdx.x, 'bid:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "        \n",
    "        dlon = y_2 - y_1\n",
    "        dlat = x_2 - x_1\n",
    "        a = sin(dlat/2)**2 + cos(x_1) * cos(x_2) * sin(dlon/2)**2\n",
    "        \n",
    "        c = 2 * asin(sqrt(a)) \n",
    "        \n",
    "        out[i] = c * r\n",
    "   </pre>\n",
    "    <pre>\n",
    "df = df.apply_rows(haversine_distance_kernel,\n",
    "                   incols=['lat1', 'lon1', 'lat2', 'lon2'],\n",
    "                   outcols=dict(out=np.float64),\n",
    "                   kwargs=dict(r=6371))\n",
    "print(df.head()\n",
    "</pre>\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#modify the code in this cell\n",
    "def haversine_distance_kernel(lat1, lon1, lat2, lon2, out, r):\n",
    "    \"\"\"Haversine distance formula taken from Michael Dunn's StackOverflow post:\n",
    "    https://stackoverflow.com/questions/4913349/haversine-formula-in-python-bearing-and-distance-between-two-gps-points\n",
    "    \"\"\"\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread_id:', cuda.threadIdx.x, 'bid:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "    x_1 = pi/180 * x_1\n",
    "    y_1 = pi/180 * y_1\n",
    "    x_2 = pi/180 * x_2\n",
    "    y_2 = pi/180 * y_2\n",
    "\n",
    "    dlon = y_2 - y_1\n",
    "    dlat = x_2 - x_1\n",
    "    a = sin(dlat/2)**2 + cos(x_1) * cos(x_2) * sin(dlon/2)**2\n",
    "\n",
    "    c = 2 * asin(sqrt(a)) \n",
    "\n",
    "    out[i] = c * r\n",
    "   \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='exercise2'></a><br>\n",
    "## **Exercise 2**\n",
    "\n",
    "Write a kernel to compute the [bearing formula](https://www.movable-type.co.uk/scripts/latlong.html):\n",
    "\n",
    "$$\\operatorname{atan2}(\\sin(\\lambda_2-\\lambda_1)\\cos(\\varphi_2), \\cos(\\varphi_1)\\sin(\\varphi_2)-\\sin(\\varphi_1)\\cos(\\varphi_2)\\cos(\\lambda_2-\\lambda_1))$$\n",
    "       \n",
    "where again $\\varphi_1,\\varphi_2$ are the latitudes and $\\lambda_1,\\lambda_2$ are the longitudes.\n",
    "\n",
    "<details><summary><b>Solution</b></summary>\n",
    "   <pre>\n",
    "from math import atan2\n",
    "\n",
    "def bearing_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread_id:', cuda.threadIdx.x, 'bid:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "            x_1 = pi/180 * x_1\n",
    "            y_1 = pi/180 * y_1\n",
    "            x_2 = pi/180 * x_2\n",
    "            y_2 = pi/180 * y_2\n",
    "\n",
    "            dlon = y_2 - y_1\n",
    "            a = atan2(sin(dlon)*cos(x_2), cos(x_1)*sin(x_1) - sin(x_1)*cos(x_2)*cos(dlon))\n",
    "            # Convert radians [-π, π] to degrees [0°, 360°]\n",
    "            out[i] = (180/pi*a + 180) % 360\n",
    "\n",
    "df = df.apply_rows(bearing_kernel,\n",
    "                   incols=['lat1', 'lon1', 'lat2', 'lon2'],\n",
    "                   outcols=dict(out=np.float64),\n",
    "                   kwargs=dict())\n",
    "\n",
    "print(df.head())\n",
    "\n",
    "</pre>\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#modify the code in this cell\n",
    "from math import atan2\n",
    "\n",
    "def bearing_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread_id:', cuda.threadIdx.x, 'bid:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "\n",
    "        dlon = y_2 - y_1\n",
    "        a = atan2(sin(dlon)*cos(x_2), cos(x_1)*sin(x_1) - sin(x_1)*cos(x_2)*cos(dlon))\n",
    "        # Convert radians [-π, π] to degrees [0°, 360°]\n",
    "        out[i] = (180/pi*a + 180) % 360"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df.head())\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='chunks'></a><br>\n",
    "\n",
    "## `apply_chunks`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the section above, the data was generally split into single-element chunks. `apply_chunks` is a more general version of `apply_rows` that gives us control over how the data is chunked on the GPU. We can specify how to divide the long array, map each of the array chunks to different GPU blocks to process (using the `chunks` argument) and assign the number of threads per block (using the `tpb` argument).\n",
    "\n",
    "Applying kernels with `apply_chunks` is very similar to applying kernels with `apply_rows`. Except, when we call `apply_chunks`, we must also provide:\n",
    "- The chunk size `chunks` as an integer or `cudf.Series` of integer offsets\n",
    "- The number of threads per block, `tpb`. Note that `tpb` can be omitted, but in that case, it defaults to `1` thread per block, which is very inefficient. We recommend always setting this argument.\n",
    "\n",
    "The kernel is executed by each thread, with full access to all the elements in that chunk of the array. In this example below, with `chunks=16`, cuDF tries to uniformly cut the 1000 elements into chunks of size 16 spread across multiple blocks. Eight threads per block process the subarray of size 16, since we set `tpb=8`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the exact same kernel as above. \n",
    "def haversine_distance_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    \"\"\"Haversine distance formula taken from Michael Dunn's StackOverflow post:\n",
    "    https://stackoverflow.com/questions/4913349/haversine-formula-in-python-bearing-and-distance-between-two-gps-points\n",
    "    \"\"\"\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "        print('thread_id:', cuda.threadIdx.x, 'bid:', cuda.blockIdx.x,\n",
    "              'array size:', lat1.size, 'block threads:', cuda.blockDim.x)\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "        \n",
    "        dlon = y_2 - y_1\n",
    "        dlat = x_2 - x_1\n",
    "        a = sin(dlat/2)**2 + cos(x_1) * cos(x_2) * sin(dlon/2)**2\n",
    "        \n",
    "        c = 2 * asin(sqrt(a)) \n",
    "        r = 6371 # Radius of earth in kilometers\n",
    "        \n",
    "        \n",
    "        out[i] = c * r\n",
    "        \n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "outdf = df.apply_chunks(haversine_distance_kernel,\n",
    "                        incols=['lat1', 'lon1', 'lat2', 'lon2'],\n",
    "                        outcols=dict(out=np.float64),\n",
    "                        kwargs=dict(),\n",
    "                        chunks=16,\n",
    "                        tpb=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(outdf.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The execution time has clearly decreased from 530 ms to 492 ms by using apply_chunks, but the emphasis is on the increased control over manner of execution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='exercise3'></a><br>\n",
    "\n",
    "## **Exercise 3**\n",
    "\n",
    "Use your bearing formula function from Exercise 2 with `apply_chunks` instead of `apply_rows` (if you did not complete the exercises, click **solution** below each exercise to get the solution).\n",
    "\n",
    "\n",
    "<details><summary><b>Solution</b></summary>\n",
    "   <pre>\n",
    "\n",
    "from math import atan2\n",
    "\n",
    "def bearing_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "\n",
    "        dlon = y_2 - y_1\n",
    "        a = atan2(sin(dlon)*cos(x_2), cos(x_1)*sin(x_1) - sin(x_1)*cos(x_2)*cos(dlon))\n",
    "        # Convert radians [-π, π] to degrees [0°, 360°]\n",
    "        out[i] = (180/pi*a + 180) % 360\n",
    "\n",
    "<br>\n",
    "\n",
    "df = df.apply_chunks(bearing_kernel,\n",
    "                     incols=['lat1', 'lon1', 'lat2', 'lon2'],\n",
    "                     outcols=dict(out=np.float64),\n",
    "                     kwargs=dict(),\n",
    "                     chunks=16,\n",
    "                     tpb=8)\n",
    "\n",
    "print(df.head())\n",
    "</pre>\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#modify the code in this cell\n",
    "from math import atan2\n",
    "\n",
    "\n",
    "def bearing_kernel(lat1, lon1, lat2, lon2, out):\n",
    "    for i, (x_1, y_1, x_2, y_2) in enumerate(zip(lat1, lon1, lat2, lon2)):\n",
    "\n",
    "\n",
    "        x_1 = pi/180 * x_1\n",
    "        y_1 = pi/180 * y_1\n",
    "        x_2 = pi/180 * x_2\n",
    "        y_2 = pi/180 * y_2\n",
    "\n",
    "        dlon = y_2 - y_1\n",
    "        a = atan2(sin(dlon)*cos(x_2), cos(x_1)*sin(x_1) - sin(x_1)*cos(x_2)*cos(dlon))\n",
    "        # Convert radians [-π, π] to degrees [0°, 360°]\n",
    "        out[i] = (180/pi*a + 180) % 360\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We can see that the computation time using apply_chunks has become 3.2 ms as compared to the 449 ms using apply_rows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='groups'></a><br>\n",
    "\n",
    "# Advanced UDFs: `apply_grouped`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we'll walk through how we can apply UDFs to a grouped DataFrame, and why you might want to do this in the first place.\n",
    "\n",
    "In the financial services industry, data scientists often need to compute features from time series data. One of the most popular ways to process time series data is to compute a moving average, as if you were sliding a window across your array. With the skills we've learned so far, you could create a custom UDF to do exactly that! You could define the function, pass it into `apply_rows` or `apply_chunks`, and get the moving average results.\n",
    "\n",
    "But, often, our DataFrame will contain *multiple* time series that we want to process logically separately (such the time series of prices for different stocks). In the introductory notebook, we learned about the `groupby` concept in cuDF that helps us define separate groups and process them separately.\n",
    "\n",
    "In the following example, we’ll show how to combine these two concepts (groupbys and UDFs) to calculate moving averages within each separate group."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll create a random 15 row DataFrame with one categorical feature and one random integer valued feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = cudf.DataFrame(\n",
    "        {\n",
    "            \"stock\": [1] * 5 + [2] * 5 + [3] * 5,\n",
    "            \"price\": [np.random.randint(0, 100) for _ in range(15)],\n",
    "        }\n",
    "     )\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll group the DataFrame by its categorical feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grouped_df = df.groupby(\"stock\", method=\"cudf\")\n",
    "grouped_df = df.groupby(\"stock\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll define a kernel which takes the moving average of a sliding window. We'll call this function `rolling_avg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rolling_avg(price, avg):\n",
    "    win_size = 3\n",
    "    for i in range(cuda.threadIdx.x, len(price), cuda.blockDim.x):\n",
    "        if i < win_size - 1:\n",
    "            # If there is not enough data to fill the window,\n",
    "            # take the average to be NaN\n",
    "            avg[i] = np.nan\n",
    "        else:\n",
    "            total = 0\n",
    "            for j in range(i - win_size + 1, i + 1):\n",
    "                total += price[j]\n",
    "            avg[i] = total / win_size\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our `rolling_avg` function defined, we can pass it to the `apply_grouped` method for grouped DataFrames and compute the moving average within each group.\n",
    "\n",
    "Currently, the function argument that corresponds to the column you're taking the average of (`price`) must match the name of your column. As a result, we name that argument `price`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "# Compute moving avgs on all groups\n",
    "results = grouped_df.apply_grouped(rolling_avg,\n",
    "                               incols=['price'],\n",
    "                               outcols=dict(avg=np.float64))\n",
    "\n",
    "end = time.time()\n",
    "print('cuDF time', end-start)\n",
    "\n",
    "print(\"Results:\")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the value for `avg` at beginning of each group is null. This makes sense, since we can't have an average value for a sliding window of size `3` until we have at least three values. We used the line `avg[i] = np.nan` in our `rolling_avg` function to geneate this result in the output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "At this point, we've introduced UDFs and the ways you can apply them to DataFrames. As we can control the data mapping to a greater extent now, the acceleration of operations is possible. Feel free to experiment writing your own UDFs in the cells below. \n",
    "For a walk-through of UDFs in more detail, we encourage you to explore our [Overview of User Defined Functions with cuDF](https://docs.rapids.ai/api/cudf/nightly/guide-to-udfs.html). After this tutorial, you will be able to handle heavy computations and data manipulation using CuDF easily. Now let's apply this knowledge to a practical exercise in the next lab. It will test your skills and help you understand how these theoretical concepts are implemented in practical data science pipelines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Licensing\n",
    "  \n",
    "This material is released by NVIDIA Corporation under the Creative Commons Attribution 4.0 International (CC BY 4.0)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Previous Notebook](01-Intro_to_cuDF.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](01-Intro_to_cuDF.ipynb)\n",
    "[2]\n",
    "[3](03-Cudf_Exercise.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[Next Notebook](03-Cudf_Exercise.ipynb)\n",
    "\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&ensp;\n",
    "[Home Page](../START_HERE.ipynb)"
   ]
  }
 ],
 "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
