improving the advanced indexing notebooks

intermediate/indexing/advanced-indexing.ipynb

@@ -8,7 +8,9 @@
     "\n",
     "## Learning Objectives\n",
     "\n",
-    "* Orthogonal vs. Vectorized and Pointwise Indexing"
+    "* Orthogonal vs. Pointwise (Vectorized) Indexing.\n",
+    "* Pointwise indexing in Xarray to extract data at a collection of points.\n",
+    "* Understand the difference between NumPy and Xarray indexing behavior."
    ]
   },
   {
@@ -17,60 +19,109 @@
    "source": [
     "## Overview\n",
     "\n",
-    "In the previous notebooks, we learned basic forms of indexing with xarray (positional and name based dimensions, integer and label based indexing), Datetime Indexing, and nearest neighbor lookups. In this tutorial, we will learn how Xarray indexing is different from Numpy and how to do vectorized/pointwise indexing using Xarray. \n",
-    "First, let's import packages needed for this repository: "
+    "In the previous notebooks, we learned basic forms of indexing with Xarray, including positional and label-based indexing, datetime indexing, and nearest neighbor lookups. We also learned that indexing an Xarray DataArray directly works (mostly) like it does for NumPy arrays; however, Xarray indexing behavior deviates from NumPy when using multiple arrays for indexing, like `arr[[0, 1], [0, 1]]`.\n",
+    "\n",
+    "To better understand this difference, let's take a look at an example of 2D 5x5 array:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "import pandas as pd\n",
-    "import xarray as xr\n",
     "\n",
+    "# Create a 5x5 array with values from 1 to 25\n",
+    "np_array = np.arange(1, 26).reshape(5, 5)\n",
+    "np_array"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now create a Xarray DataArray from this NumPy array: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
     "\n",
-    "xr.set_options(display_expand_attrs=False)\n",
-    "np.set_printoptions(threshold=10, edgeitems=2)"
+    "da = xr.DataArray(np_array, dims=[\"x\", \"y\"])\n",
+    "da"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In this notebook, we’ll use air temperature tutorial dataset from the National Center for Environmental Prediction.  "
+    "Now, let's see how the indexing behavior is different between NumPy array and Xarray DataArray when indexing with multiple arrays:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "ds = xr.tutorial.load_dataset(\"air_temperature\")\n",
-    "da = ds.air\n",
-    "ds"
+    "np_array[[0, 2, 4], [0, 2, 4]]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da[[0, 2, 4], [0, 2, 4]]"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Orthogonal Indexing \n",
+    "The image below summarizes the difference between vectorized and orthogonal indexing for a 2D 5x5 NumPy array and Xarray DataArray:\n",
+    "\n",
+    "\n",
     "\n",
-    "As we learned in the previous tutorial, positional indexing deviates from the behavior exhibited by NumPy when indexing with multiple arrays. However, Xarray pointwise indexing supports the indexing along multiple labeled dimensions using list-like objects similar to NumPy indexing behavior.\n",
+    "![Orthogonal vs. Vectorized Indexing](../../images/orthogonal_vs_vectorized.png)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Pointwise** or **Vectorized indexing**, shown on the left, selects specific elements at given coordinates, resulting in an array of those individual elements. In the example shown, the indices `[0, 2, 4]`, `[0, 2, 4]` select the elements at positions `(0, 0)`, `(2, 2)`, and `(4, 4)`, resulting in the values `[1, 13, 25]`. This is the default behavior of NumPy arrays.\n",
+    "  \n",
+    "In contrast, **orthogonal indexing** uses the same indices to select entire rows and columns, forming a cross-product of the specified indices. This method results in sub-arrays that include all combinations of the selected rows and columns. The example demonstrates this by selecting rows 0, 2, and 4 and columns 0, 2, and 4, resulting in a subarray containing `[[1, 3, 5], [11, 13, 15], [21, 23, 25]]`. This is Xarray DataArray's default behavior.\n",
+    " \n",
+    "The output of vectorized indexing is a `1D array`,  while the output of orthogonal indexing is a `3x3` array.   \n",
     "\n",
-    "If you only provide integers, slices, or unlabeled arrays (array without dimension names, such as `np.ndarray`, `list`, but not `DataArray()`) indexing can be understood as orthogonally (i.e. along independent axes, instead of using NumPy’s broadcasting rules to vectorize indexers). \n",
     "\n",
-    "*Orthogonal* or *outer* indexing considers one-dimensional arrays in the same way as slices when deciding the output shapes. The principle of outer or orthogonal indexing is that the result mirrors the effect of independently indexing along each dimension with integer or boolean arrays, treating both the indexed and indexing arrays as one-dimensional. This method of indexing is analogous to vector indexing in programming languages like MATLAB, Fortran, and R, where each indexer component *independently* selects along its corresponding dimension. \n",
+    ":::{tip}  To Summarize: \n",
+    "\n",
+    "- *Pointwise* or *vectorized* indexing is a more general form of indexing that allows for arbitrary combinations of indexing arrays. This method of indexing is analogous to the broadcasting rules in NumPy, where the dimensions of the indexers are aligned and the result is determined by the shape of the indexers. This is the default behavior in NumPy.\n",
+    "\n",
+    "- *Orthogonal* or *outer* indexing allows for indexing along each dimension independently, treating the indexers as one-dimensional arrays. The principle of outer or orthogonal indexing is that the result mirrors the effect of independently indexing along each dimension with integer or boolean arrays, treating both the indexed and indexing arrays as one-dimensional. This method of indexing is analogous to vector indexing in programming languages like MATLAB, Fortran, and R, where each indexer component independently selects along its corresponding dimension. This is the default behavior in Xarray.\n",
+    "\n",
+    "\n",
+    ":::"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Orthogonal Indexing in Xarray\n",
+    "\n",
+    "As explained earlier, when you use only integers, slices, or unlabeled arrays (arrays without dimension names, such as `np.ndarray` or `list`, but not `DataArray`) to index an `Xarray DataArray`, Xarray interprets these indexers orthogonally. This means it indexes along independent axes, rather than using NumPy's broadcasting rules to vectorize the indexers. \n",
     "\n",
-    "For example : "
+    "In the example above we saw this behavior, but let's see this behavior in action with a real dataset. Here we’ll use `air temperature` data from the National Center for Environmental Prediction:"
    ]
   },
   {
@@ -81,14 +132,40 @@
    },
    "outputs": [],
    "source": [
-    "da.isel(time=0, lat=[2, 4, 10, 13], lon=[1, 6, 7]).plot();  # -- orthogonal indexing"
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "\n",
+    "\n",
+    "xr.set_options(display_expand_attrs=False)\n",
+    "np.set_printoptions(threshold=10, edgeitems=2)\n",
+    "\n",
+    "ds = xr.tutorial.load_dataset(\"air_temperature\")\n",
+    "da_air = ds.air\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selected_da = da_air.isel(time=0, lat=[2, 4, 10, 13], lon=[1, 6, 7])  # -- orthogonal indexing\n",
+    "selected_da"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "For more flexibility, you can supply `DataArray()` objects as indexers. Dimensions on resultant arrays are given by the ordered union of the indexers’ dimensions:\n",
+    "👆 please notice how the output of the indexing example above resulted in an array of size `3x4`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For more flexibility, you can supply `DataArray()` objects as indexers. Dimensions on resultant arrays are given by the ordered union of the indexers’ dimensions.\n",
     "\n",
     "For example, in the example below we do orthogonal indexing using `DataArray()` objects. "
    ]
@@ -104,14 +181,30 @@
     "target_lat = xr.DataArray([31, 41, 42, 42], dims=\"degrees_north\")\n",
     "target_lon = xr.DataArray([200, 201, 202, 205], dims=\"degrees_east\")\n",
     "\n",
-    "da.sel(lat=target_lat, lon=target_lon, method=\"nearest\")  # -- orthogonal indexing"
+    "da_air.sel(lat=target_lat, lon=target_lon, method=\"nearest\")  # -- orthogonal indexing"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In the above example, you can see how the output shape is `time` x `lats` x `lons`. "
+    "In the above example, you can see how the output shape is `time` x `lats` x `lons`. Please note that there are no shared dimensions between the indexers, so the output shape is the union of the dimensions of the indexers.\n",
+    "\n",
+    "```{attention}\n",
+    "Please note that slices or sequences/arrays without named-dimensions are treated as if they have the same dimension which is indexed along.\n",
+    "```\n",
+    "\n",
+    "For example:\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_air.sel(lat=[20, 30, 40], lon=target_lon, method=\"nearest\")"
    ]
   },
   {
@@ -121,20 +214,54 @@
    },
    "source": [
     "\n",
-    "But what if we would like to find the information from the nearest grid cell to a collection of specified points (for example, weather stations or tower data)?\n",
+    "But what if we'd like to find the nearest climate model grid cell to a collection of specified points (for example observation sites, or weather stations)?\n",
     "\n",
-    "## Vectorized or Pointwise Indexing\n",
+    "## Vectorized or Pointwise Indexing in Xarray\n",
     "\n",
-    "Like NumPy and pandas, Xarray supports indexing many array elements at once in a\n",
-    "*vectorized* manner. \n",
+    "Like NumPy and pandas, Xarray supports indexing many array elements at once in a *vectorized* manner. \n",
     "\n",
-    "**Vectorized indexing** or **Pointwise Indexing** using `DataArrays()` can be used to extract information from the nearest grid cells of interest, for example, the nearest climate model grid cells to a collection of specified weather station latitudes and longitudes.\n",
+    "**Vectorized indexing** or **Pointwise Indexing** using `DataArrays()` can be used to extract information from the nearest grid cells of interest, for example, the nearest climate model grid cells to a collection of specified observation tower data latitudes and longitudes.\n",
     "\n",
     "```{hint}\n",
-    "To trigger vectorized indexing behavior, you will need to provide the selection dimensions with a new shared output dimension name. \n",
+    "To trigger vectorized indexing behavior, you will need to provide the selection dimensions with a new **shared** output dimension name. This means that the dimensions of both indexers must be the same, and the output will have the same dimension name as the indexers.\n",
     "```\n",
     "\n",
-    "In the example below, the selections of the closest latitude and longitude are renamed to an output dimension named `points`:"
+    "Let's see how this works with an example. A researcher wants to find the nearest climate model grid cell to a collection of observation sites. She has the latitude and longitude of the observation sites as following:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "obs_lats = [31.81, 41.26, 22.59, 44.47, 28.57]\n",
+    "\n",
+    "obs_lons = [200.16, 201.57, 305.54, 210.56, 226.59]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If the researcher use the lists to index the DataArray, they will get the orthogonal indexing behavior, which is not what they wants."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_air.sel(lat=obs_lats, lon=obs_lats, method=\"nearest\")  # -- orthogonal indexing"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To trigger the pointwise indexing, they need to create DataArray objects with the same dimension name, and then use them to index the DataArray. \n",
+    "For example, the code below first create DataArray objects for the latitude and longitude of the observation sites using a shared dimension name `points`, and then use them to index the DataArray `air_temperature`:"
    ]
   },
   {
@@ -145,9 +272,8 @@
    },
    "outputs": [],
    "source": [
-    "# Define target latitude and longitude (where weather stations might be)\n",
-    "lat_points = xr.DataArray([31, 41, 42, 42], dims=\"points\")\n",
-    "lon_points = xr.DataArray([200, 201, 202, 205], dims=\"points\")\n",
+    "## latitudes of weather stations with a dimension of \"points\"\n",
+    "lat_points = xr.DataArray(obs_lats, dims=\"points\")\n",
     "lat_points"
    ]
   },
@@ -159,6 +285,8 @@
    },
    "outputs": [],
    "source": [
+    "## longitudes of weather stations with a dimension of \"points\"\n",
+    "lon_points = xr.DataArray(obs_lons, dims=\"points\")\n",
     "lon_points"
    ]
   },
@@ -177,7 +305,7 @@
    },
    "outputs": [],
    "source": [
-    "da.sel(lat=lat_points, lon=lon_points, method=\"nearest\")"
+    "da_air.sel(lat=lat_points, lon=lon_points, method=\"nearest\")  # -- pointwise indexing"
    ]
   },
   {
@@ -195,38 +323,50 @@
    },
    "outputs": [],
    "source": [
-    "da.sel(lat=lat_points, lon=lon_points, method=\"nearest\").dims"
+    "da_air.sel(lat=lat_points, lon=lon_points, method=\"nearest\").dims"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "```{attention}\n",
-    "Please note that slices or sequences/arrays without named-dimensions are treated as if they have the same dimension which is indexed along.\n",
-    "```\n",
-    "\n",
-    "For example:"
+    "Now, let's plot the data for all stations."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "da.sel(lat=[20, 30, 40], lon=lon_points, method=\"nearest\")"
+    "da_air.sel(lat=lat_points, lon=lon_points, method=\"nearest\").plot(x='time', hue='points');"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "```{warning}\n",
-    "If an indexer is a `DataArray()`, its coordinates should not conflict with the selected subpart of the target array (except for the explicitly indexed dimensions with `.loc`/`.sel`). Otherwise, `IndexError` will be raised!\n",
-    "```"
+    "## Exercises\n",
+    "\n",
+    "::::{admonition} Exercise\n",
+    ":class: tip\n",
+    "\n",
+    "In the simple 2D 5x5 Xarray data array above, select the sub-array containing (0,0),(2,2),(4,4):\n",
+    "\n",
+    ":::{admonition} Solution\n",
+    ":class: dropdown\n",
+    "```python\n",
+    "\n",
+    "indices = np.array([0, 2, 4])\n",
+    "\n",
+    "xs_da = xr.DataArray(indices, dims=\"points\")\n",
+    "ys_da = xr.DataArray(indices, dims=\"points\")\n",
+    "\n",
+    "subset_da = da.sel(x=xs_da, y=xs_da)\n",
+    "subset_da\n",
+    "```\n",
+    ":::\n",
+    "::::"
    ]
   },
   {
@@ -235,7 +375,14 @@
    "source": [
     "## Additional Resources\n",
     "\n",
-    "- [Xarray Docs - Indexing and Selecting Data](https://docs.xarray.dev/en/stable/indexing.html)\n"
+    "- [Xarray Docs - Indexing and Selecting Data](https://docs.xarray.dev/en/stable/indexing.html)\n",
+    "\n",
+    "\n",
+    ":::{seealso}\n",
+    "- [Introductions to Fancy Indexing](https://jakevdp.github.io/PythonDataScienceHandbook/02.07-fancy-indexing.html)\n",
+    "- [NumPy Docs - Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing)\n",
+    "\n",
+    ":::\n"
    ]
   }
  ],