docs/aggregation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aaf2f2cb",
   "metadata": {},
   "source": [
    "# Aggregation\n",
    "\n",
    "Xarray-Beam can perform efficient distributed data aggregation in the \"map-reduce\" model. \n",
    "\n",
    "This currently only includes `Mean`, but we would welcome contributions of other aggregation functions such as `Sum`, `Std`, `Var`, `Min`, `Max`, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b77b953",
   "metadata": {},
   "source": [
    "## High-level API\n",
    "\n",
    "The `Mean` transformation comes in three forms:  {py:class}`Mean <xarray_beam.Mean>`, {py:class}`Mean.Globally <xarray_beam.Mean.Globally>`, and {py:class}`Mean.PerKey <xarray_beam.Mean.PerKey>`. The implementation is highly scalable, based on a Beam's [`CombineFn`](https://beam.apache.org/documentation/transforms/python/aggregation/combineglobally/#example-4-combining-with-a-combinefn)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca45196",
   "metadata": {},
   "source": [
    "The high-level `Mean` transform can be used to aggregate a distributed dataset across an existing dimension or dimensions, similar to Xarray's `.mean()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3e387dd1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset>\n",
      "Dimensions:  (lat: 25, time: 2920, lon: 53)\n",
      "Coordinates:\n",
      "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
      "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
      "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
      "Data variables:\n",
      "    air      (time, lat, lon) float32 241.2 242.5 243.5 ... 296.5 296.2 295.7\n",
      "Attributes:\n",
      "    Conventions:  COARDS\n",
      "    title:        4x daily NMC reanalysis (1948)\n",
      "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
      "    platform:     Model\n",
      "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...\n"
     ]
    }
   ],
   "source": [
    "import apache_beam as beam\n",
    "import numpy as np\n",
    "import xarray_beam as xbeam\n",
    "import xarray\n",
    "\n",
    "ds = xarray.tutorial.load_dataset('air_temperature')\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5967340a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "        if (typeof window.interactive_beam_jquery == 'undefined') {\n",
       "          var jqueryScript = document.createElement('script');\n",
       "          jqueryScript.src = 'https://code.jquery.com/jquery-3.4.1.slim.min.js';\n",
       "          jqueryScript.type = 'text/javascript';\n",
       "          jqueryScript.onload = function() {\n",
       "            var datatableScript = document.createElement('script');\n",
       "            datatableScript.src = 'https://cdn.datatables.net/1.10.20/js/jquery.dataTables.min.js';\n",
       "            datatableScript.type = 'text/javascript';\n",
       "            datatableScript.onload = function() {\n",
       "              window.interactive_beam_jquery = jQuery.noConflict(true);\n",
       "              window.interactive_beam_jquery(document).ready(function($){\n",
       "                \n",
       "              });\n",
       "            }\n",
       "            document.head.appendChild(datatableScript);\n",
       "          };\n",
       "          document.head.appendChild(jqueryScript);\n",
       "        } else {\n",
       "          window.interactive_beam_jquery(document).ready(function($){\n",
       "            \n",
       "          });\n",
       "        }"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[2]: Mean/PerKey/CombinePerKey(MeanCombineFn)/GroupByKey'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[2]: Mean/PerKey/CombinePerKey(MeanCombineFn)/GroupByKey'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[2]: Mean/PerKey/CombinePerKey(MeanCombineFn)'. \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(Key(offsets={'lat': 0, 'lon': 0}, vars=None), <xarray.Dataset>\n",
      "Dimensions:  (lat: 25, lon: 53)\n",
      "Coordinates:\n",
      "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
      "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
      "Data variables:\n",
      "    air      (lat, lon) float64 260.4 260.2 259.9 259.5 ... 297.3 297.3 297.3)\n"
     ]
    }
   ],
   "source": [
    "with beam.Pipeline() as p:\n",
    "    p | xbeam.DatasetToChunks(ds, chunks={'time': 1000}) | xbeam.Mean('time') | beam.Map(print)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a584485",
   "metadata": {},
   "source": [
    "## Lower-level API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7585a9",
   "metadata": {},
   "source": [
    "Xarray-Beam also includes lower-level transforations modelled off of [`beam.Mean`](https://beam.apache.org/documentation/transforms/python/aggregation/mean/) rather than {py:meth}`xarray.Dataset.mean`: they compute averages over sequences of `xarray.Dataset` objects or (`key`, `xarray.Dataset`) pairs, rather than calculating an average over an existing Xarray dimension or based on `xarray_beam.Key` objects, e.g.,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "86a925af",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['/Users/shoyer/miniconda3/envs/xarray-beam/lib/python3.9/site-packages/ipykernel_launcher.py', '-f', '/Users/shoyer/Library/Jupyter/runtime/kernel-b19331c2-e263-4737-bd64-012081154884.json']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<xarray.Dataset>\n",
       " Dimensions:  (x: 3)\n",
       " Dimensions without coordinates: x\n",
       " Data variables:\n",
       "     foo      (x) float64 0.05667 -0.02306 -0.1648]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets = [\n",
    "    xarray.Dataset({'foo': ('x', np.random.randn(3))})\n",
    "    for _ in range(100)\n",
    "]\n",
    "datasets | xbeam.Mean.Globally()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5670130b",
   "metadata": {},
   "source": [
    "Notice how existing dimensions on each datasets are unchanged by the transformation. If you want to average over existing dimensions, use the high-level `Mean` transform or do that aggregation yourself, e.g., by averaging inside each chunk before combining the data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab152b6",
   "metadata": {},
   "source": [
    "Similarly, the keys fed into `xbeam.Mean.PerKey` can be any hashables, including but not limited to `xbeam.Key`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b2399483",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['/Users/shoyer/miniconda3/envs/xarray-beam/lib/python3.9/site-packages/ipykernel_launcher.py', '-f', '/Users/shoyer/Library/Jupyter/runtime/kernel-b19331c2-e263-4737-bd64-012081154884.json']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[('DJF',\n",
       "  <xarray.Dataset>\n",
       "  Dimensions:  ()\n",
       "  Data variables:\n",
       "      air      float64 273.6),\n",
       " ('MAM',\n",
       "  <xarray.Dataset>\n",
       "  Dimensions:  ()\n",
       "  Data variables:\n",
       "      air      float64 279.0),\n",
       " ('JJA',\n",
       "  <xarray.Dataset>\n",
       "  Dimensions:  ()\n",
       "  Data variables:\n",
       "      air      float64 289.2),\n",
       " ('SON',\n",
       "  <xarray.Dataset>\n",
       "  Dimensions:  ()\n",
       "  Data variables:\n",
       "      air      float64 283.0)]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets = [\n",
    "    (time.dt.season.item(), ds.sel(time=time).mean())\n",
    "    for time in ds.time\n",
    "]\n",
    "datasets | xbeam.Mean.PerKey()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9db1627",
   "metadata": {},
   "source": [
    "`Mean.PerKey` is particularly useful in combination with {class}`beam.GroupByKey` for performing large-scale \"group by\" operations. For example, that a look at the [ERA5 climatology example](https://github.com/google/xarray-beam/blob/main/examples/era5_climatology.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb6fc6a",
   "metadata": {},
   "source": [
    "## Custom aggregations\n",
    "\n",
    "The \"tree reduction\" algorithm used by the combiner inside `Mean` is great, but it isn't the only way to aggregate a dataset with Xarray-Beam.\n",
    "\n",
    "In many cases, the easiest way to scale up an aggregation pipeline is to make use of [rechunking](rechunking.ipynb) to convert the many small datasets inside your pipeline into a form that is easier to calculate in a scalable way. However, rechunking is much less efficient than using combiner, because each use of `Rechunk` requires a complete shuffle of the input data (i.e., writing all data in the pipepilne to temporary files on disk).\n",
    "\n",
    "For example, here's how one could compute the `median`, which is a notoriously difficult statistic to calculate with distributed algorithms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ef1ef099",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: Rechunk/Stage0/Consolidate/GroupByTempKeys'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: Rechunk/Stage0/Consolidate/GroupByTempKeys'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: Rechunk/Stage2/Consolidate/GroupByTempKeys'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: Rechunk/Stage2/Consolidate/GroupByTempKeys'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: ConsolidateChunks/GroupByTempKeys'. \n",
      "WARNING:apache_beam.coders.coder_impl:Using fallback deterministic coder for type '<class 'xarray_beam._src.core.Key'>' in '[5]: ConsolidateChunks/GroupByTempKeys'. \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset>\n",
      "Dimensions:  (lat: 25, lon: 53)\n",
      "Coordinates:\n",
      "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
      "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
      "Data variables:\n",
      "    air      (lat, lon) float32 261.3 261.1 260.9 260.3 ... 297.3 297.3 297.3\n"
     ]
    }
   ],
   "source": [
    "source_chunks = {'time': 100, 'lat': -1, 'lon': -1}\n",
    "working_chunks = {'lat': 10, 'lon': 10, 'time': -1}\n",
    "\n",
    "with beam.Pipeline() as p:\n",
    "    (\n",
    "        p\n",
    "        | xbeam.DatasetToChunks(ds, source_chunks)\n",
    "        | xbeam.Rechunk(ds.sizes, source_chunks, working_chunks, itemsize=4)\n",
    "        | beam.MapTuple(lambda k, v: (k.with_offsets(time=None), v.median('time')))\n",
    "        | xbeam.ConsolidateChunks({'lat': -1, 'lon': -1})\n",
    "        | beam.MapTuple(lambda k, v: print(v))\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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}