2-custom-operations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f3e6c3f1-9673-4eb0-81b1-3fe9ac7e84cf",
   "metadata": {},
   "source": [
    "<img src=\"images/dask_horizontal.svg\" width=\"45%\" alt=\"Dask logo\\\">\n",
    "     \n",
    "# Custom Operations\n",
    "\n",
    "In the overview notebook, we discussed some of the many algorithms that are pre-defined for different types of Dask collections\n",
    "(such as Arrays and DataFrames).\n",
    "These include operations like `mean`, `max`, `value_counts`, and many other standard operations.\n",
    "\n",
    "In this notebook we'll:\n",
    " - explore how those operations are implemented\n",
    " - learn how to construct our own custom operations and optimizations\n",
    " - get a sense of when you might need to write a custom operations\n",
    "\n",
    "**Related Documentation**\n",
    "\n",
    "  - [Array Tutorial](https://tutorial.dask.org/02_array.html)\n",
    "  - [Best Practices for Customization](https://docs.dask.org/en/latest/best-practices.html#learn-techniques-for-customization)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c0fe241-96cf-414d-b5ee-d700b51651df",
   "metadata": {},
   "source": [
    "## Blocked Algorithms\n",
    "\n",
    "Dask computations are implemented using _blocked algorithms_. These algorithms break up a computation on a large array into many computations on smaller pieces of the array. This minimizes the memory load (amount of RAM) of computations and allows for working with larger-than-memory datasets in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca494a4-362a-4503-9725-e79ebeee84c4",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask.array as da\n",
    "\n",
    "x = da.random.random(size=(1_000, 1_000), chunks=(250, 500))\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "849c5b30-fab9-4cd8-9e7f-a2dd4bbf1f4c",
   "metadata": {},
   "source": [
    "In the overview notebook we looked at the task graph for the following computation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "776c7645-3c80-4390-a8ed-dde0d4463712",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = (x + x.T).sum(axis=0).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9f181e4-0205-4f9a-89db-4f8f3a54d801",
   "metadata": {},
   "source": [
    "Now let's break that down a bit and look at the task graph for just one part of that computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad44e637-1782-4eb1-89c3-0025230e1fce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x.T.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4087e539-2d92-4ed4-98ab-03471de15c8d",
   "metadata": {},
   "source": [
    "This graph demonstrates how blocked algorithms work. In the perfectly parallelizable situation, Dask can operate on each block in isolation and then reassemble the results from the outputs. Dask makes it easy to contruct graphs like this using a numpy-like API. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "negative-minority",
   "metadata": {},
   "source": [
    "## Custom Block Computations\n",
    "Block computations operate on a per-block basis. So each block gets the function applied to it, and the output has the same chunk location as the input.\n",
    "\n",
    "Some examples include the following:\n",
    "- custom IO operations\n",
    "- applying embarassingly parallel functions for which there is no exising Dask implementation\n",
    "\n",
    "![map_blocks](images/custom_operations_map_blocks.png)\n",
    "\n",
    "**Related Documentation**\n",
    "\n",
    "   - [`dask.array.map_blocks`](https://docs.dask.org/en/latest/generated/dask.array.map_blocks.html#dask.array.map_blocks)\n",
    "   - [`dask.dataframe.map_partitions`](https://dask.pydata.org/en/latest/generated/dask.dataframe.DataFrame.map_partitions.html#dask.dataframe.DataFrame.map_partitions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wired-albert",
   "metadata": {},
   "source": [
    "### `map_blocks`\n",
    "\n",
    "Let's imagine that there was no `da.random.random` method. We can create our own version using `map_blocks`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8262fd43-b432-4e40-aced-4f9ad633d0c7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def random_sample():\n",
    "    return np.random.random(size=(250, 500))\n",
    "\n",
    "x = da.map_blocks(random_sample, chunks=((250, 250, 250, 250), (500, 500)), dtype=float)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8660b3ed-494e-4c41-af0f-59b0e1626194",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "requested-congo",
   "metadata": {},
   "source": [
    "#### Understanding `chunks` argument\n",
    "\n",
    "In the example above we explicitly declare what the size of the output chunks will be ``chunks=((250, 250, 250, 250), (500, 500))`` this means 8 chunks each with shape `(250, 500)` you'll also see the chunks argument written in the short version where only the shape of one chunk is defined ``chunks=(250, 500)``. These mean the same thing.\n",
    "\n",
    "Specifying the output chunks is very useful when doing more involved operations with ``map_blocks``. By specifying ``chunks``, you can guarantee that the output will have the right shape. Having the right shape lets you properly chain together other operations. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "555f0551-bf38-4e95-a8a4-3aa9ac5b26fd",
   "metadata": {},
   "source": [
    "In that example we created an array from scratch by passing in `dtype` and `chunks`. Next we'll consider the case of applying `map_blocks` to existing arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nervous-devon",
   "metadata": {},
   "source": [
    "#### Multiple arrays\n",
    "\n",
    "``map_blocks`` can be used on single arrays or to combine several arrays. When multiple arrays are passed, ``map_blocks``\n",
    "aligns blocks by block location without regard to shape.\n",
    "\n",
    "In the following example we have two arrays with the same number of blocks\n",
    "but with different shape and chunk sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "effective-correspondence",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = da.arange(1000, chunks=(100,))\n",
    "b = da.arange(100, chunks=(10,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "biblical-maldives",
   "metadata": {},
   "source": [
    "Let's take a look at these arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "representative-twins",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sound-johns",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "polish-excess",
   "metadata": {},
   "source": [
    "We can pass these arrays into ``map_blocks`` using a function that takes two inputs, calculates the max of each, then returns a numpy array of the outputs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exceptional-aberdeen",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def func(a, b):\n",
    "    return np.array([a.max(), b.max()])\n",
    "\n",
    "result = da.map_blocks(func, a, b, chunks=(2,))\n",
    "result.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adapted-output",
   "metadata": {},
   "source": [
    "#### Special arguments\n",
    "\n",
    "There are special arguments (``block_info`` and ``block_id``) that you can use within ``map_blocks`` functions. \n",
    "\n",
    " - ``block_id`` gives the index of the block within the chunks, so for a 1D array it will be something like `(i,)`. \n",
    " - ``block_info`` is a dictionary where there is an integer key for each input dask array and a `None` key for the output array.\n",
    " \n",
    "These special arguments let you know where you are in the ``map_blocks`` from within the custom function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "atlantic-interim",
   "metadata": {},
   "source": [
    "### ``map_partitions``\n",
    "\n",
    "In Dask dataframe there is a similar method to ``map_blocks`` but it is called ``map_partitions``.\n",
    "\n",
    "Here is an example of using it to check if the sum of two columns is greater than some arbitrary threshold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01aeac1-a313-4ab4-b940-719ee5fbd116",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "import dask.dataframe as dd\n",
    "\n",
    "ddf = dask.datasets.timeseries()\n",
    "\n",
    "result = ddf.map_partitions(lambda df, threshold: (df.x + df.y) > 0, threshold=0)\n",
    "result.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "collected-preparation",
   "metadata": {},
   "source": [
    "#### Internal uses\n",
    "\n",
    "In practice ``map_partitions`` is used to implement many of the helper dataframe methods\n",
    "that let Dask dataframe mimic Pandas. Here is the [implementation](https://github.com/dask/dask/blob/67e648922512615f94f8a90726423e721d0e3eb2/dask/dataframe/core.py#L662-L671) of `ddf.index` for instance:\n",
    "\n",
    "```python\n",
    "@property\n",
    "def index(self):\n",
    "    \"\"\"Return dask Index instance\"\"\"\n",
    "    return self.map_partitions(\n",
    "        getattr,\n",
    "        \"index\",\n",
    "        token=key_split(self._name) + \"-index\",\n",
    "        meta=self._meta.index,\n",
    "        enforce_metadata=False,\n",
    "    )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "optical-norwegian",
   "metadata": {},
   "source": [
    "#### Understanding `meta` argument\n",
    "\n",
    "Dask dataframes and dask arrays have a special attribute called `_meta` that allows them to know metadata about the type of dataframe/array that they represent. This metadata includes:\n",
    " - dtype (int, float)\n",
    " - column names and order\n",
    " - name\n",
    " - type (pandas dataframe, cudf dataframe)\n",
    " \n",
    "**Related documentation**\n",
    "\n",
    "- [Dataframe metadata](https://docs.dask.org/en/latest/dataframe-design.html#metadata)\n",
    "\n",
    "This information is stored in an empty object of the proper type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "external-treaty",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(ddf._meta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tender-architect",
   "metadata": {},
   "source": [
    "That's how Dask knows what to render when you display a dask object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "operating-estimate",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(ddf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conscious-prairie",
   "metadata": {},
   "source": [
    "When you add an item to the task graph, Dask tries to run the function on the meta before you call compute. \n",
    "\n",
    "This approach has several benefits:\n",
    "\n",
    "- it gives Dask a sense of what the output will look like. \n",
    "- if there are fundamental issues, Dask will fail fast\n",
    "\n",
    "Here's a few examples. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd3fb018-de20-4f75-9d7d-441a6ec3a97b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "essential-excitement",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.name.str.startswith(\"A\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcfa14a6-fd76-4dff-89de-def4039b0103",
   "metadata": {},
   "source": [
    "See how the output looks right? The dtypes are correct, the type is a `Series` rather than a `DataFrame` like the input.\n",
    "\n",
    "**Exercise**\n",
    "\n",
    "Try using `startswith` on a different column and see what you get :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7344e1d3-e417-4324-82e0-57e8efab5fd4",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# solution\n",
    "ddf.x.str.startswith(\"A\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "congressional-memory",
   "metadata": {},
   "source": [
    "### Declaring meta\n",
    "\n",
    "When setting up custom operations, sometimes running the function on a miniature version of the data doesn't produce a result that is similar enough to your expected output.\n",
    "\n",
    "In those cases you can manually provide a `meta` to use as the `meta` of the output. You're basically telling Dask what to expect from the output. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "royal-carroll",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = ddf.map_partitions(lambda df, threshold: (df.x + df.y) > threshold, threshold=0, meta=(\"greater_than\", bool))\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4627e5ae-f31a-4868-85d4-3b8119870c49",
   "metadata": {},
   "source": [
    "Note that Dask is trusting you to give it the correct meta. There is no enforcement."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expressed-sharing",
   "metadata": {},
   "source": [
    "### `map_overlap`\n",
    "Sometimes you want to operate on a per-block basis, but you need some information from neighboring blocks. \n",
    "\n",
    "Example operations include the following:\n",
    "\n",
    "- Convolve a filter across an image\n",
    "- Rolling sum/mean/max, …\n",
    "- Search for image motifs like a Gaussian blob that might span the border of a block\n",
    "- Evaluate a partial derivative\n",
    "\n",
    "Dask Array supports these operations by creating a new array where each block is slightly expanded by the borders of its neighbors. \n",
    "\n",
    "![](https://docs.dask.org/en/latest/_images/overlapping-neighbors.svg)\n",
    "\n",
    "This costs an excess copy and the communication of many small chunks, but allows localized functions to evaluate in an embarrassingly parallel manner.\n",
    "\n",
    "**Related Documentation**\n",
    "   - [Array Overlap](https://docs.dask.org/en/latest/array-overlap.html)\n",
    "\n",
    "The main API for these computations is the ``map_overlap`` method. ``map_overlap`` is very similar to ``map_blocks`` but has the additional arguments: ``depth``, ``boundary``, and ``trim``.\n",
    "\n",
    "Here is an example of calculating the derivative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e78c226-a318-4772-959d-f099f8564185",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import dask.array as da\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "a = np.array([1, 1, 2, 3, 3, 3, 2, 1, 1])\n",
    "a = da.from_array(a, chunks=5)\n",
    "\n",
    "plt.plot(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cde80b97-1781-4104-8ca2-b43e4544fc4c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def derivative(a):\n",
    "    return a - np.roll(a, 1)\n",
    "\n",
    "b = a.map_overlap(derivative, depth=1, boundary=None)\n",
    "b.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "restricted-lying",
   "metadata": {},
   "source": [
    "In this case each block shares 1 value from its neighboring block: ``depth``. And since we set ``boundary=0``on the outer edges of the array, the first and last block are padded with the integer 0. Since we haven't specified ``trim`` it is true by default meaning that the overlap is removed before returning the results.\n",
    "\n",
    "If you inspect the task graph you'll see two mostly independent towers of tasks, with just some value sharing at the edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eight-mirror",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "b.visualize(collapse_outputs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "great-crowd",
   "metadata": {},
   "source": [
    "**Exercise**\n",
    "\n",
    "Lets apply a Gaussian filter to an image following the example from the [SciPy docs](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.gaussian_filter.html).\n",
    "\n",
    "First create a Dask array from the NumPy array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "handmade-infrared",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy.datasets import ascent\n",
    "import dask.array as da\n",
    "\n",
    "a = da.from_array(ascent(), chunks=(128, 128))\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "disabled-editing",
   "metadata": {},
   "source": [
    "Now use ``map_overlap`` to apply ``gaussian_filter`` to each block.\n",
    "\n",
    "```python\n",
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "b = a.map_overlap(gaussian_filter, sigma=5, ...)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "informative-worcester",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# solution\n",
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "b = a.map_overlap(gaussian_filter, sigma=5, depth=10, boundary=\"periodic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "printable-anaheim",
   "metadata": {},
   "source": [
    "Check what you've come up with by plotting the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "satisfactory-basket",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))\n",
    "ax1.imshow(a)\n",
    "ax2.imshow(b)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2908a29e-061e-4964-9018-c1ea1240ba42",
   "metadata": {},
   "source": [
    "Notice that if you set the depth to a smaller value, you can see the edges of the blocks in the output image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22cb0c3f-1aca-42db-9ec3-565a45c68a45",
   "metadata": {},
   "source": [
    "## Reduction\n",
    "Each dask collection has a `reduction` method. This is the generalized method that supports operations that reduce the dimensionality of the inputs.\n",
    "\n",
    "![Custom operations: reduction](images/custom_operations_reduction.png)\n",
    "\n",
    "**Related Documentation**\n",
    "   - [`dask.array.reduction`](https://dask.pydata.org/en/latest/generated/dask.array.reduction.html#dask.array.reduction)\n",
    "   - [`dask.dataframe.reduction`](https://dask.pydata.org/en/latest/generated/dask.dataframe.DataFrame.reduction.html#dask.dataframe.DataFrame.reduction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "unexpected-representation",
   "metadata": {},
   "source": [
    "### Internal uses\n",
    "\n",
    "This is the [internal definition](https://github.com/dask/dask/blob/67e648922512615f94f8a90726423e721d0e3eb2/dask/array/reductions.py#L394-L407) of `sum` on `dask.Array`. In it you can see that there is a\n",
    "regular `np.sum` applied across each block and then tree-reduced with `np.sum` again.\n",
    "\n",
    "```python\n",
    "def sum(a, axis=None, dtype=None, keepdims=False, split_every=None, out=None):\n",
    "    if dtype is None:\n",
    "        dtype = getattr(np.zeros(1, dtype=a.dtype).sum(), \"dtype\", object)\n",
    "    result = reduction(\n",
    "        a,\n",
    "        chunk.sum,\n",
    "        chunk.sum,\n",
    "        axis=axis,\n",
    "        keepdims=keepdims,\n",
    "        dtype=dtype,\n",
    "        split_every=split_every,\n",
    "        out=out,\n",
    "    )\n",
    "    return result\n",
    "```\n",
    "\n",
    "Here is `da.sum` reimplemented as a custom reduction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eab1d5a-c839-4b7c-8fed-9bf88bfa12d5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "da.reduction(x, np.sum, np.sum, dtype=x.dtype).visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "negative-display",
   "metadata": {},
   "source": [
    "By visualizing `b`, we can see how the tree reduction works. First, `sum` is applied to each block, then every 4 chunks are combined using `sum-partial`.\n",
    "This keeps going until there are less than 4 results left, then `sum-aggregate` is used to finish up."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "greek-heritage",
   "metadata": {},
   "source": [
    "**Exercise**\n",
    "\n",
    "See how the graph changes when you set the chunks - maybe to `(100, 250)` or `(250, 250)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7a74a2-4e02-4807-b4d1-2221818ef95d",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# solution\n",
    "x = da.random.random(size=(1_000, 1_000), chunks=(100, 250))\n",
    "x.sum().visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb0f3153-ac83-46c3-abaa-3e487a81c0b4",
   "metadata": {},
   "source": [
    "### Understanding `split_every`\n",
    "\n",
    "`split_every` controls the number of chunk outputs that are used as input to each `partial` call. \n",
    "\n",
    "Here is an example of doing partial aggregation on every 5 blocks along the 0 axis and every 2 blocks along the 1 axis (so 10 blocks go into each `partial-sum`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e57621ca-58c4-4506-964f-9cf53746e962",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x.sum(split_every={0: 5, 1: 2}).visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "902b2b79-af4d-4d79-bd2e-daa90e80bcd5",
   "metadata": {},
   "source": [
    "**Exercise**\n",
    "\n",
    "Try setting different values for `split_every` and visualizing the task graph to see the impact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b1dcaa-83fa-474f-9351-dc443500595b",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# solution\n",
    "x.sum(split_every={0: 10, 1: 2}).visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e5bd9c4-5aaf-4e09-b655-210bf0cc3169",
   "metadata": {},
   "source": [
    "> **Side note**\n",
    ">\n",
    "> You can use reductions to calculate aggregations per-block reduction even if you don't want to combine and aggregate the results of those blocks:\n",
    ">\n",
    "> ```python\n",
    "> da.reduction(x, np.sum, lambda x, **kwargs: x, dtype=int).compute()\n",
    "> ```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "disturbed-snapshot",
   "metadata": {},
   "source": [
    "## When to use which method\n",
    "\n",
    "In this notebook we've covered several different mechanisms for applying arbitrary functions to the blocks in arrays or dataframes. Here's a brief summary of when you should use these various methods\n",
    "\n",
    "- `map_block`, `map_partition` - block organization of the input matches the block organization of the output and the function is fully parallelizable. \n",
    "- `map_overlap` - block organizations of input and output match, but the function is not fully parallelizable (requires input from neighboring chunks).\n",
    "- `blockwise` - same function can be applied to the blocks as to the partial and aggregated versions. Also output blocks can be in different orientations.\n",
    "- `reduction` - dimensionality of output does not necessarily match that of input and function is fully parallelizable.\n",
    "- `groupby().agg` - data needs to be aggregated per group (the index of the output will be the group keys).\n",
    "- `dask.delayed` - data doesn't have a complex block organization or the data is small and the computation is pretty fast."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "stone-surfing",
   "metadata": {},
   "source": [
    "## Customizing Optimization\n",
    "\n",
    "Dask defines a default optimization strategy for each collection type (Array,\n",
    "Bag, DataFrame, Delayed). However, different applications may have different\n",
    "needs.  \n",
    "\n",
    "Oftentimes this looks like turning off particular types of optimization using \n",
    "the Dask config settings:\n",
    "\n",
    "```python\n",
    "dask.config.set({\"optimization.fuse.active\": False})\n",
    "```\n",
    "\n",
    "You can also construct your own custom optimization function and use it \n",
    "instead of the default. An optimization function takes in a task graph \n",
    "and list of desired keys and returns a new task graph:\n",
    "\n",
    "```python\n",
    "def my_optimize_function(dsk, keys):\n",
    "    new_dsk = {...}\n",
    "    return new_dsk\n",
    "```\n",
    "\n",
    "You can then register this optimization class against whichever collection type\n",
    "you prefer and it will be used instead of the default scheme:\n",
    "\n",
    "```python\n",
    "with dask.config.set(array_optimize=my_optimize_function):\n",
    "    x, y = dask.compute(x, y)\n",
    "```\n",
    "\n",
    "You can register separate optimization functions for different collections, or\n",
    "you can register ``None`` if you do not want particular types of collections to\n",
    "be optimized:\n",
    "\n",
    "```python\n",
    "with dask.config.set(array_optimize=my_optimize_function,\n",
    "                     dataframe_optimize=None,\n",
    "                     delayed_optimize=my_other_optimize_function):\n",
    "    ...\n",
    "```\n",
    "\n",
    "You do not need to specify all collections.  Collections will default to their\n",
    "standard optimization scheme (which is usually a good choice).\n",
    "\n",
    "# Next steps\n",
    "\n",
    "Next, let's learn more about Dask's distributed scheduler in the [3-distributed-scheduler.ipynb](3-distributed-scheduler.ipynb) notebook."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all",
   "text_representation": {
    "extension": ".md",
    "format_name": "markdown"
   }
  },
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}