Add multidimensional binning demo

Author: dcherian

docs/source/user-stories/nD-bins.ipynb

@@ -0,0 +1,339 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e970d800-c612-482a-bb3a-b1eb7ad53d88",
+   "metadata": {},
+   "source": [
+    "# Multi-dimensional Binning\n",
+    "\n",
+    "```{warning}\n",
+    "This post is a proof-of-concept for discussion. Expect APIs to change to enable this use case.\n",
+    "```\n",
+    "\n",
+    "Here we explore a binning problem where the bins are multidimensional\n",
+    "([xhistogram issue](https://github.com/xgcm/xhistogram/issues/28))\n",
+    "\n",
+    "> One of such multi-dim bin applications is the ranked probability score rps we\n",
+    "> use in `xskillscore.rps`, where we want to know how many forecasts fell into\n",
+    "> which bins. Bins are often defined as terciles of the forecast distribution\n",
+    "> and the bins for these terciles\n",
+    "> (`forecast_with_lon_lat_time_dims.quantile(q=[.33,.66],dim='time')`) depend on\n",
+    "> `lon` and `lat`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01f1a2ef-de62-45d0-a04e-343cd78debc5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "import flox\n",
+    "import flox.xarray\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import xarray as xr"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0be3e214-0cf0-426f-8ebb-669cc5322310",
+   "metadata": {},
+   "source": [
+    "## Create test data\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ce239000-e053-4fc3-ad14-e9e0160da869",
+   "metadata": {},
+   "source": [
+    "Data to be reduced\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7659c24e-f5a1-4e59-84c0-5ec965ef92d2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "array = xr.DataArray(\n",
+    "    np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),\n",
+    "    dims=(\"space\", \"time\"),\n",
+    "    name=\"array\",\n",
+    ")\n",
+    "array"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "da0c0ac9-ad75-42cd-a1ea-99069f5bef00",
+   "metadata": {},
+   "source": [
+    "Array to group by\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4601e744-5d22-447e-97ce-9644198d485e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "by = xr.DataArray(\n",
+    "    np.array([[1, 2, 3], [3, 4, 5], [5, 6, 7], [6, 7, 9]]),\n",
+    "    dims=(\"space\", \"time\"),\n",
+    "    name=\"by\",\n",
+    ")\n",
+    "by"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "61c21c94-7b6e-46a6-b9c2-59d7b2d40c81",
+   "metadata": {},
+   "source": [
+    "multidimensional bins:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "863a1991-ab8d-47c0-aa48-22b422fcea8c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bins = by + 0.5\n",
+    "bins = xr.DataArray(\n",
+    "    np.concatenate([bins, bins[:, [-1]] + 1], axis=-1)[:, :-1].T,\n",
+    "    dims=(\"time\", \"nbins\"),\n",
+    ")\n",
+    "# array += xr.DataArray([0, 1, 2, 4], dims=\"space\")\n",
+    "\n",
+    "display(by)\n",
+    "display(bins)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e65ecaba-d1cc-4485-ae58-c390cb2ebfab",
+   "metadata": {},
+   "source": [
+    "## Concept\n",
+    "\n",
+    "The key idea is that GroupBy is two steps:\n",
+    "\n",
+    "1. Factorize (a.k.a \"digitize\") : convert the `by` data to a set of integer\n",
+    "   codes representing the bins.\n",
+    "2. Apply the reduction.\n",
+    "\n",
+    "We treat multi-dimensional binning as a slightly complicated factorization\n",
+    "problem. Assume that bins are a function of `time`. So we\n",
+    "\n",
+    "1. generate a set of appropriate integer codes by:\n",
+    "   1. Loop over \"time\" and factorize the data appropriately.\n",
+    "   2. Add an offset to these codes so that \"bin 0\" for `time=0` is different\n",
+    "      from \"bin 0\" for `time=1`\n",
+    "2. apply the groupby reduction to the \"offset codes\"\n",
+    "3. reshape the output to the right shape\n",
+    "\n",
+    "We will work at the xarray level, so its easy to keep track of the different\n",
+    "dimensions.\n",
+    "\n",
+    "### Factorizing\n",
+    "\n",
+    "The core `factorize_` function (which wraps `pd.cut`) only handles 1D bins, so\n",
+    "we use `xr.apply_ufunc` to vectorize it for us.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa33ab2c-0ecf-4198-a033-2a77f5d83c99",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "factorize_loop_dim = \"time\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "afcddcc1-dd57-461e-a649-1f8bcd30342f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def factorize_nd_bins_core(by, bins):\n",
+    "    group_idx, *_, props = flox.core.factorize_(\n",
+    "        (by,),\n",
+    "        axis=-1,\n",
+    "        expected_groups=(pd.IntervalIndex.from_breaks(bins),),\n",
+    "    )\n",
+    "    # Use -1 as the NaN sentinel value\n",
+    "    group_idx[props.nanmask] = -1\n",
+    "    return group_idx\n",
+    "\n",
+    "\n",
+    "codes = xr.apply_ufunc(\n",
+    "    factorize_nd_bins_core,\n",
+    "    by,\n",
+    "    bins,\n",
+    "    # TODO: avoid hardcoded dim names\n",
+    "    input_core_dims=[[\"space\"], [\"nbins\"]],\n",
+    "    output_core_dims=[[\"space\"]],\n",
+    "    vectorize=True,\n",
+    ")\n",
+    "codes"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1661312a-dc61-4a26-bfd8-12c2dc01eb15",
+   "metadata": {},
+   "source": [
+    "### Offset the codes\n",
+    "\n",
+    "These are integer codes appropriate for a single timestep.\n",
+    "\n",
+    "We now add an offset that changes in time, to make sure \"bin 0\" for `time=0` is\n",
+    "different from \"bin 0\" for `time=1` (taken from\n",
+    "[this StackOverflow thread](https://stackoverflow.com/questions/46256279/bin-elements-per-row-vectorized-2d-bincount-for-numpy)).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e5801cb-a79c-4670-ad10-36bb19f1a6ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N = math.prod([codes.sizes[d] for d in codes.dims if d != factorize_loop_dim])\n",
+    "offset = xr.DataArray(\n",
+    "    np.arange(codes.sizes[factorize_loop_dim]), dims=factorize_loop_dim\n",
+    ")\n",
+    "# TODO: think about N-1 here\n",
+    "offset_codes = (codes + offset * (N - 1)).rename(by.name)\n",
+    "offset_codes.data[codes == -1] = -1\n",
+    "offset_codes"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c06c48b-316b-4a33-9bc3-921acd10bcba",
+   "metadata": {},
+   "source": [
+    "### Reduce\n",
+    "\n",
+    "Now that we have appropriate codes, let's apply the reduction\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2cf1295e-4585-48b9-ac2b-9e00d03b2b9a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "interim = flox.xarray.xarray_reduce(\n",
+    "    array,\n",
+    "    offset_codes,\n",
+    "    func=\"sum\",\n",
+    "    # We use RangeIndex to indicate that `-1` code can be safely ignored\n",
+    "    # (it indicates values outside the bins)\n",
+    "    # TODO: Avoid hardcoding 9 = sizes[\"time\"] x (sizes[\"nbins\"] - 1)\n",
+    "    expected_groups=pd.RangeIndex(9),\n",
+    ")\n",
+    "interim"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3539509b-d9b4-4342-a679-6ada6f285dfb",
+   "metadata": {},
+   "source": [
+    "## Make final result\n",
+    "\n",
+    "Now reshape that 1D result appropriately.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1389d37-d76d-4a50-9dfb-8710258de3fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "final = (\n",
+    "    interim.coarsen(by=3)\n",
+    "    # bin_number dimension is last, this makes sense since it is the core dimension\n",
+    "    # and we vectorize over the loop dims.\n",
+    "    # So the first (Nbins-1) elements are for the first index of the loop dim\n",
+    "    .construct({\"by\": (factorize_loop_dim, \"bin_number\")})\n",
+    "    .transpose(..., factorize_loop_dim)\n",
+    "    .drop_vars(\"by\")\n",
+    ")\n",
+    "final"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a98b5e60-94af-45ae-be1b-4cb47e2d77ba",
+   "metadata": {},
+   "source": [
+    "I think this is the expected answer.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "053a8643-f6d9-4fd1-b014-230fa716449c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "array.isel(space=slice(1, None)).rename({\"space\": \"bin_number\"}).identical(\n",
+    "    final\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "619ba4c4-7c87-459a-ab86-c187d3a86c67",
+   "metadata": {},
+   "source": [
+    "## TODO\n",
+    "\n",
+    "This could be extended to:\n",
+    "\n",
+    "1. handle multiple `factorize_loop_dim`\n",
+    "2. avoid hard coded dimension names in the `apply_ufunc` call for factorizing\n",
+    "3. avoid hard coded number of output elements in the `xarray_reduce` call.\n",
+    "4. Somehow propagate the bin edges to the final output.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "flox-tests",
+   "language": "python",
+   "name": "conda-env-flox-tests-py"
+  },
+  "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.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}