From 09de2887339ac334a0eca3524a428410dcacde8a Mon Sep 17 00:00:00 2001 From: dcherian Date: Tue, 17 Jan 2023 21:49:18 -0700 Subject: [PATCH 1/3] Add nD binning notebook --- docs/source/user-stories.md | 1 + docs/source/user-stories/nD-bins.ipynb | 330 +++++++++++++++++++++++++ 2 files changed, 331 insertions(+) create mode 100644 docs/source/user-stories/nD-bins.ipynb diff --git a/docs/source/user-stories.md b/docs/source/user-stories.md index 22b37939e..0241e01dc 100644 --- a/docs/source/user-stories.md +++ b/docs/source/user-stories.md @@ -8,4 +8,5 @@ user-stories/climatology.ipynb user-stories/climatology-hourly.ipynb user-stories/custom-aggregations.ipynb + user-stories/nD-bins.ipynb ``` diff --git a/docs/source/user-stories/nD-bins.ipynb b/docs/source/user-stories/nD-bins.ipynb new file mode 100644 index 000000000..eb97fa21b --- /dev/null +++ b/docs/source/user-stories/nD-bins.ipynb @@ -0,0 +1,330 @@ +{ + "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 numpy as np\n", + "import pandas as pd\n", + "import xarray as xr\n", + "\n", + "import flox\n", + "import flox.xarray" + ] + }, + { + "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(np.arange(codes.sizes[factorize_loop_dim]), dims=factorize_loop_dim)\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(final)" + ] + }, + { + "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": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 096d2994700fcf8d03776541c1b5294a7d7d3806 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Wed, 4 Oct 2023 21:48:59 -0600 Subject: [PATCH 2/3] fix. --- docs/source/user-stories/nD-bins.ipynb | 97 +++++++++++++++++++------- 1 file changed, 70 insertions(+), 27 deletions(-) diff --git a/docs/source/user-stories/nD-bins.ipynb b/docs/source/user-stories/nD-bins.ipynb index eb97fa21b..52948e41f 100644 --- a/docs/source/user-stories/nD-bins.ipynb +++ b/docs/source/user-stories/nD-bins.ipynb @@ -3,7 +3,10 @@ { "cell_type": "markdown", "id": "e970d800-c612-482a-bb3a-b1eb7ad53d88", - "metadata": {}, + "metadata": { + "tags": [], + "user_expressions": [] + }, "source": [ "# Multi-dimensional Binning\n", "\n", @@ -26,7 +29,9 @@ "cell_type": "code", "execution_count": null, "id": "01f1a2ef-de62-45d0-a04e-343cd78debc5", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "import math\n", @@ -42,7 +47,9 @@ { "cell_type": "markdown", "id": "0be3e214-0cf0-426f-8ebb-669cc5322310", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "## Create test data\n" ] @@ -50,7 +57,9 @@ { "cell_type": "markdown", "id": "ce239000-e053-4fc3-ad14-e9e0160da869", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "Data to be reduced\n" ] @@ -59,7 +68,9 @@ "cell_type": "code", "execution_count": null, "id": "7659c24e-f5a1-4e59-84c0-5ec965ef92d2", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "array = xr.DataArray(\n", @@ -73,7 +84,9 @@ { "cell_type": "markdown", "id": "da0c0ac9-ad75-42cd-a1ea-99069f5bef00", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "Array to group by\n" ] @@ -82,7 +95,9 @@ "cell_type": "code", "execution_count": null, "id": "4601e744-5d22-447e-97ce-9644198d485e", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "by = xr.DataArray(\n", @@ -96,33 +111,38 @@ { "cell_type": "markdown", "id": "61c21c94-7b6e-46a6-b9c2-59d7b2d40c81", - "metadata": {}, + "metadata": { + "tags": [], + "user_expressions": [] + }, "source": [ - "multidimensional bins:\n" + "Multidimensional bins:\n" ] }, { "cell_type": "code", "execution_count": null, "id": "863a1991-ab8d-47c0-aa48-22b422fcea8c", - "metadata": {}, + "metadata": { + "tags": [] + }, "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", + " name=\"bins\",\n", ")\n", - "# array += xr.DataArray([0, 1, 2, 4], dims=\"space\")\n", - "\n", - "display(by)\n", - "display(bins)" + "bins" ] }, { "cell_type": "markdown", "id": "e65ecaba-d1cc-4485-ae58-c390cb2ebfab", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "## Concept\n", "\n", @@ -155,7 +175,9 @@ "cell_type": "code", "execution_count": null, "id": "aa33ab2c-0ecf-4198-a033-2a77f5d83c99", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "factorize_loop_dim = \"time\"" @@ -165,13 +187,15 @@ "cell_type": "code", "execution_count": null, "id": "afcddcc1-dd57-461e-a649-1f8bcd30342f", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "def factorize_nd_bins_core(by, bins):\n", " group_idx, *_, props = flox.core.factorize_(\n", " (by,),\n", - " axis=-1,\n", + " axes=(-1,),\n", " expected_groups=(pd.IntervalIndex.from_breaks(bins),),\n", " )\n", " # Use -1 as the NaN sentinel value\n", @@ -194,7 +218,9 @@ { "cell_type": "markdown", "id": "1661312a-dc61-4a26-bfd8-12c2dc01eb15", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "### Offset the codes\n", "\n", @@ -209,7 +235,9 @@ "cell_type": "code", "execution_count": null, "id": "0e5801cb-a79c-4670-ad10-36bb19f1a6ff", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "N = math.prod([codes.sizes[d] for d in codes.dims if d != factorize_loop_dim])\n", @@ -223,7 +251,9 @@ { "cell_type": "markdown", "id": "6c06c48b-316b-4a33-9bc3-921acd10bcba", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "### Reduce\n", "\n", @@ -234,7 +264,9 @@ "cell_type": "code", "execution_count": null, "id": "2cf1295e-4585-48b9-ac2b-9e00d03b2b9a", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "interim = flox.xarray.xarray_reduce(\n", @@ -252,7 +284,9 @@ { "cell_type": "markdown", "id": "3539509b-d9b4-4342-a679-6ada6f285dfb", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "## Make final result\n", "\n", @@ -263,7 +297,9 @@ "cell_type": "code", "execution_count": null, "id": "b1389d37-d76d-4a50-9dfb-8710258de3fd", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "final = (\n", @@ -281,7 +317,9 @@ { "cell_type": "markdown", "id": "a98b5e60-94af-45ae-be1b-4cb47e2d77ba", - "metadata": {}, + "metadata": { + "user_expressions": [] + }, "source": [ "I think this is the expected answer.\n" ] @@ -290,7 +328,9 @@ "cell_type": "code", "execution_count": null, "id": "053a8643-f6d9-4fd1-b014-230fa716449c", - "metadata": {}, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ "array.isel(space=slice(1, None)).rename({\"space\": \"bin_number\"}).identical(final)" @@ -299,7 +339,10 @@ { "cell_type": "markdown", "id": "619ba4c4-7c87-459a-ab86-c187d3a86c67", - "metadata": {}, + "metadata": { + "tags": [], + "user_expressions": [] + }, "source": [ "## TODO\n", "\n", From b905277b37947fd58035d4863d36757a7e9e68f2 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Wed, 4 Oct 2023 21:50:01 -0600 Subject: [PATCH 3/3] title --- docs/source/user-stories/nD-bins.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/user-stories/nD-bins.ipynb b/docs/source/user-stories/nD-bins.ipynb index 52948e41f..87ef942bf 100644 --- a/docs/source/user-stories/nD-bins.ipynb +++ b/docs/source/user-stories/nD-bins.ipynb @@ -8,7 +8,7 @@ "user_expressions": [] }, "source": [ - "# Multi-dimensional Binning\n", + "# Binning with multi-dimensional bins\n", "\n", "```{warning}\n", "This post is a proof-of-concept for discussion. Expect APIs to change to enable this use case.\n",