{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93434031-d7fe-4322-a9cf-41e5b8be622d",
   "metadata": {},
   "source": [
    "# High-level Computation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c0164c9-f970-4206-964a-d1d8080e2b0a",
   "metadata": {},
   "source": [
    "## Overview\n",
    "### In this tutorial, you learn:\n",
    "\n",
    "* High level Xarray computations using CuPy arrays. \n",
    "* Applying custom kernels to DataArray with CuPy and NumPy\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "| Concepts | Importance | Notes |\n",
    "| --- | --- | --- |\n",
    "| [Familiarity with NumPy](https://foundations.projectpythia.org/core/numpy.html) | Necessary | |\n",
    "| [Basics of Cupy](Notebook0_Introduction) | Necessary | |\n",
    "| [Familiarity with Xarray](https://foundations.projectpythia.org/core/xarray.html) | Necessary | |\n",
    "\n",
    "- **Time to learn**: 40 minutes\n",
    "\n",
    "\n",
    "\n",
    "## Introduction \n",
    "\n",
    "In the previous tutorial, we introduced the powerful combination of Xarray and CuPy for handling multi-dimensional datasets and leveraging GPU acceleration to significantly improve performance. \n",
    "\n",
    "In this tutorial, we are going to explore high-level Xarray functions such as `groupby`, `rolling mean`, and `weighted mean`, and compared their execution times with traditional NumPy-based implementations.\n",
    "\n",
    "##  High-level Xarray Functions: CuPy vs. NumPy\n",
    "\n",
    "In this tutorial, we'll explore the performance differences between high-level Xarray functions using CuPy and NumPy. CuPy is a GPU-based NumPy-compatible library, while NumPy is the well-known CPU-based library for numerical computations. We'll focus on three high-level functions: groupby, rolling mean, and weighted mean. We'll also compare the time it takes to execute each function using both CuPy and NumPy.\n",
    "Let's create some sample data to work with.\n",
    "\n",
    "We'll use a 3-dimensional dataset (time, latitude, longitude) with random values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e21d74da-e620-4737-91a4-0180f3703c75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "377ac7ca-6c09-486f-95cc-8a2d599df800",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "447116b1-bc54-4d0a-af46-5908a7f95e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "import cupy_xarray  # Adds .cupy to Xarray objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5f6f83a3-36d6-493d-bfda-49f5d766ef14",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "\n",
    "# Create the time range.\n",
    "date = pd.date_range(\"2010-01-01\", \"2020-12-31\", freq=\"M\")\n",
    "\n",
    "# Create the latitude range.\n",
    "lat = np.arange(-90, 90, 1)\n",
    "\n",
    "# Create the longitude range.\n",
    "lon = np.arange(-180, 180, 1)\n",
    "\n",
    "# Create random data\n",
    "data_np = np.random.rand(len(date), len(lat), len(lon))\n",
    "data_cp = cp.array(data_np)\n",
    "\n",
    "# -- Create DataArray with Numpy data\n",
    "data_xr_np = xr.DataArray(\n",
    "    data_np,\n",
    "    dims=[\"time\", \"lat\", \"lon\"],\n",
    "    coords=[date, lat, lon],\n",
    ")\n",
    "\n",
    "# -- Create DataArray with CuPy data\n",
    "data_xr_cp = xr.DataArray(\n",
    "    data_cp,\n",
    "    dims=[\"time\", \"lat\", \"lon\"],\n",
    "    coords=[date, lat, lon],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ea80193-172e-45c9-af18-1ec288c4bbd8",
   "metadata": {},
   "source": [
    "### Groupby\n",
    "The `groupby` function is used to group data based on one or more dimensions. Here, we'll group our data by the season in the `time` dimension using both CuPy and NumPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e1c0b119-23f9-4397-957a-ace2a607ab6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_np = time.time()\n",
    "\n",
    "grouped_data_np = data_xr_np.groupby(\"time.season\")\n",
    "mean_np = grouped_data_np.mean()\n",
    "\n",
    "end_time_np = time.time()\n",
    "time_np = end_time_np - start_time_np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd95f64d-87cc-4ae5-95f5-2170db1eb547",
   "metadata": {},
   "source": [
    "The data type of data in grouped_data_np is `numpy.ndarray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "53419533-005b-4ceb-9ae5-2deaa316c52a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[type(arr.data) for group, arr in grouped_data_np]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "973e4b7e-1b2f-4fbb-a4da-a7a47b1be327",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GroupBy with Xarray DataArrays using CuPy provides a 79.83 x speedup over NumPy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "start_time_cp = time.time()\n",
    "\n",
    "grouped_data_cp = data_xr_cp.groupby(\"time.season\")\n",
    "mean_cp = grouped_data_cp.mean()\n",
    "\n",
    "end_time_cp = time.time()\n",
    "time_cp = end_time_cp - start_time_cp\n",
    "\n",
    "print(\n",
    "    \"GroupBy with Xarray DataArrays using CuPy provides a\",\n",
    "    round(time_np / time_cp, 2),\n",
    "    \"x speedup over NumPy.\\n\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d28001bb-f360-4b50-84c6-aec922807a24",
   "metadata": {},
   "source": [
    "What about the CuPy arrays? Does it preserve the array type?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eadab73b-769a-4a43-800e-dde016e6834c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[cupy.ndarray, cupy.ndarray, cupy.ndarray, cupy.ndarray]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[type(arr.data) for group, arr in grouped_data_cp]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc3c4f8c-2077-4cdd-b86d-3a6643e38a4d",
   "metadata": {},
   "source": [
    "### What about different sizes of arrays? How does the performance compare then?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e13ecad-4b00-4766-afc3-91a9525eed7b",
   "metadata": {},
   "source": [
    "The example above showed a 1 degree DataArray. What if we increase the data size to 0.5 degree?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "24f7fdbe-4c71-40b3-84af-564fd987f96b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the latitude range.\n",
    "lat = np.arange(-90, 90, 0.5)\n",
    "\n",
    "# Create the longitude range.\n",
    "lon = np.arange(-180, 180, 0.5)\n",
    "\n",
    "# Create random data\n",
    "data_np = np.random.rand(len(date), len(lat), len(lon))\n",
    "data_cp = cp.array(data_np)\n",
    "\n",
    "# -- Create DataArray with Numpy data\n",
    "data_xr_np = xr.DataArray(\n",
    "    data_np,\n",
    "    dims=[\"time\", \"lat\", \"lon\"],\n",
    "    coords=[date, lat, lon],\n",
    ")\n",
    "\n",
    "# -- Create DataArray with CuPy data\n",
    "data_xr_cp = xr.DataArray(\n",
    "    data_cp,\n",
    "    dims=[\"time\", \"lat\", \"lon\"],\n",
    "    coords=[date, lat, lon],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0403170b-f1b2-4883-8945-af16e5571496",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_np = time.time()\n",
    "\n",
    "grouped_data_np = data_xr_np.groupby(\"time.season\").mean()\n",
    "mean_np = grouped_data_np.mean()\n",
    "\n",
    "end_time_np = time.time()\n",
    "time_np = end_time_np - start_time_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d6adf806-83ff-4b1f-954c-2a7d4ac9e0c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GroupBy with Xarray DataArrays using CuPy provides a 89.87 x speedup over NumPy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "start_time_cp = time.time()\n",
    "\n",
    "grouped_data_cp = data_xr_cp.groupby(\"time.season\").mean()\n",
    "mean_cp = grouped_data_cp.mean()\n",
    "\n",
    "end_time_cp = time.time()\n",
    "time_cp = end_time_cp - start_time_cp\n",
    "\n",
    "print(\n",
    "    \"GroupBy with Xarray DataArrays using CuPy provides a\",\n",
    "    round(time_np / time_cp, 2),\n",
    "    \"x speedup over NumPy.\\n\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bc0b970-7893-4a9c-b06a-941c134b8c2d",
   "metadata": {},
   "source": [
    "```{attention}\n",
    "Is this consistent with what you have learned in the previous modules? What if we test a very low resolution dataset?\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3d10f6d-74a3-494d-a564-a780415373b2",
   "metadata": {},
   "source": [
    "### Rolling Mean\n",
    "\n",
    "The `rolling()` method is available in DataArray objects, providing support for rolling window aggregation. This feature allows for the computation of aggregated values over a sliding window of data points within the array.\n",
    "\n",
    "A rolling window refers to a fixed-size window that moves sequentially across the data, calculating aggregated statistics or applying functions to the values within each window. This capability is particularly useful for analyzing time series or spatial data, where examining data within a specific window can reveal patterns, trends, or relationships.\n",
    "\n",
    "The rolling mean is a widely used technique for smoothing data over a specified window. \n",
    "\n",
    "In the example below, we calculate the rolling mean along the `'time'` dimension with a window size of 10:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d30fc60b-d947-43ec-8699-852c1aca6b59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.core.options.set_options at 0x2ad1d0ccdac0>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.set_options(use_bottleneck=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a0202ef1-e8a4-4676-9b20-3e91a175b78f",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_np = time.time()\n",
    "\n",
    "rolling_mean_np = data_xr_np.rolling(time=10).mean()\n",
    "\n",
    "end_time_np = time.time()\n",
    "time_np = end_time_np - start_time_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2f496410-b347-4f20-a1de-44e39129177d",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_cp = time.time()\n",
    "\n",
    "rolling_mean_cp = data_xr_cp.rolling(time=10).mean()\n",
    "\n",
    "end_time_cp = time.time()\n",
    "time_cp = end_time_cp - start_time_cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1211e89c-7d17-462f-99ac-18a71245adf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rolling mean with Xarray DataArrays using CuPy provides a 30.22 x speedup over NumPy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Rolling mean with Xarray DataArrays using CuPy provides a\",\n",
    "    round(time_np / time_cp, 2),\n",
    "    \"x speedup over NumPy.\\n\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75ff2c6-1f65-4545-bf76-274ff1ec3939",
   "metadata": {},
   "source": [
    "### Weighted Array Reductions\n",
    "\n",
    "Weighted array reductions in Xarray empower users with the ability to perform aggregations on multidimensional arrays while considering the weights assigned to each element. They currently support weighted `sum`, `mean`, `std`, `var` and `quantile`. By default, aggregation results in Xarray's rolling window operations are assigned the coordinate at the end of each window. However, it is possible to center the results by specifying `center=True` when creating the Rolling object. \n",
    "\n",
    "For example, the weighted mean is another way to smooth data, taking into account the varying importance of each data point. \n",
    "\n",
    "Here, we'll use a uniform weight along the `time` dimension:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1a580f4f-769c-46fe-84e7-1a8372362ad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_np = time.time()\n",
    "\n",
    "weights_np = xr.DataArray(np.ones_like(data_np), dims=[\"time\", \"lat\", \"lon\"])\n",
    "weighted_mean_np = data_xr_np.weighted(weights_np).mean(dim=\"time\")\n",
    "\n",
    "end_time_np = time.time()\n",
    "time_np = end_time_np - start_time_np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "bcce9e70-67c9-4228-ac18-311fa3a555f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_cp = time.time()\n",
    "\n",
    "weights_cp = xr.DataArray(cp.ones_like(data_cp), dims=[\"time\", \"lat\", \"lon\"])\n",
    "weighted_mean_cp = data_xr_cp.weighted(weights_cp).mean(dim=\"time\")\n",
    "\n",
    "end_time_cp = time.time()\n",
    "time_cp = end_time_cp - start_time_cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c48a20d7-5d83-433e-8d8f-86aba0c5cf4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weighted mean with Xarray DataArrays using CuPy provides a 13.32 x speedup over NumPy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Weighted mean with Xarray DataArrays using CuPy provides a\",\n",
    "    round(time_np / time_cp, 2),\n",
    "    \"x speedup over NumPy.\\n\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16cd108a-0bea-42a3-aa59-41acd505f5f6",
   "metadata": {},
   "source": [
    "Similarly we can calculate weighted sum or weighted quantile, etc. To learn more about weighted array reduction, please see [the user guide](https://docs.xarray.dev/en/stable/user-guide/computation.html#weighted-array-reductions)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5ce7661-da08-4f51-8f3c-a67cc8dacd83",
   "metadata": {},
   "source": [
    "### Coarsen large arrays\n",
    "\n",
    "In Xarray, the `coarsen` operation is a powerful tool for downsampling or reducing the size of large arrays. When dealing with large datasets, coarsening allows for efficient summarization of data by aggregating multiple values into a single value within a defined coarsening window. This process is particularly useful when working with high-resolution or fine-grained data, as it enables the transformation of large arrays into smaller ones while preserving the overall structure and key characteristics of the data. \n",
    "\n",
    "In order to take a block mean for every 3 days along time dimension and every 2 points along lat and lon, we can use the following: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "f9198d05-ece4-4383-8b7e-2598efba49af",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_np = time.time()\n",
    "\n",
    "coarsen_np = data_xr_np.coarsen(time=3, lat=2, lon=2).mean()\n",
    "\n",
    "end_time_np = time.time()\n",
    "time_np = end_time_np - start_time_np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f6cc471-8e3e-4471-8ffc-2f90de161055",
   "metadata": {},
   "source": [
    "`coarsen` also works in similar fashion when using CuPy:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "67bdfb24-90ac-4502-a57e-c3406381988b",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time_cp = time.time()\n",
    "\n",
    "coarsen_cp = data_xr_cp.coarsen(time=3, lat=2, lon=2).mean()\n",
    "\n",
    "end_time_cp = time.time()\n",
    "time_cp = end_time_cp - start_time_cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "9a41faf8-74f4-4a33-bf7d-c77797e8e784",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coarsen with Xarray DataArrays using CuPy provides a 443.79 x speedup over NumPy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Coarsen with Xarray DataArrays using CuPy provides a\",\n",
    "    round(time_np / time_cp, 2),\n",
    "    \"x speedup over NumPy.\\n\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95dccf03-bb64-43ef-883e-9b84081c4514",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we have learned about:\n",
    "    \n",
    "* High level Xarray computations using CuPy arrays. \n",
    "* Applying custom kernels to DataArray with CuPy and NumPy\n",
    "\n",
    "```{seealso}\n",
    "[CuPy User Guide](https://docs.cupy.dev/en/stable/user_guide/index.html)  \n",
    "[Xarray User Guide](https://docs.xarray.dev/en/stable/user-guide/index.html)  \n",
    "[Cupy-Xarray Github](https://github.com/xarray-contrib/cupy-xarray.git)  \n",
    "```"
   ]
  }
 ],
 "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.10.9"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}