4-Application-Dask-Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f94e474",
   "metadata": {},
   "source": [
    "<center>\n",
    "    \n",
    "# 4 - Combining `Dask.array` and `Numba`\n",
    "\n",
    "   \n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ed88c2",
   "metadata": {},
   "source": [
    "### Let's see some useful ways ot combine the power of `Dask` and `Numba` together!\n",
    "Some interesting links: \n",
    "- https://examples.dask.org/applications/stencils-with-numba.html\n",
    "- https://developer.nvidia.com/blog/accelerated-portfolio-construction-with-numba-and-dask-in-python/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54421968",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import dask.array as da\n",
    "import dask_image.ndfilters\n",
    "import numba\n",
    "import numpy as np\n",
    "\n",
    "from utils import show_images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a31b721",
   "metadata": {},
   "source": [
    "## Let's load some big data! \n",
    "We are going to **lazy-load** some images acquired by the [Hubble Space Telescope](https://www.wikiwand.com/en/Hubble_Space_Telescope). \n",
    "\n",
    "This image is known as the ***Hubble Ultra Deep Field***, and captures a view of nearly 10,000 galaxies (is the deepest visible-light image of the cosmos). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e0a5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "hubble_image = da.from_zarr(os.path.join(\"imgs\", \"hubble_enh.zarr\"))\n",
    "print(f\"The image has {hubble_image[..., 0].size / 1e6 : .0f} MPix, and takes:\")\n",
    "hubble_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f9f2295",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Plot a few chunks\n",
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        hubble_image.blocks[0, 23],\n",
    "        hubble_image.blocks[4, 19],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22a77686",
   "metadata": {},
   "source": [
    "Imagine that you are an astronomer and want to find all forming galaxies in their early stages (red, circular and small).\n",
    "\n",
    "Here we will use Dask and Numba to do so in an efficient manner:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e53771a",
   "metadata": {},
   "source": [
    "## Application: The Structure Tensor    \n",
    "\n",
    "The [Structure Tensor](https://www.wikiwand.com/en/Structure_tensor) is a powerful tool for analyzing the structure of images and extracting useful information from them, often used in\n",
    "- image segmentation, \n",
    "- object recognition, and\n",
    "- optical flow estimation\n",
    "\n",
    "It makes a great showcase example for *Numba* and *Dask* as it makes very **simple** operations (convolutions) and it is **highly parallelizable**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c417561",
   "metadata": {},
   "source": [
    "<div style=\"text-align: right\"><a href=\"https://www.crisluengo.net/archives/1132/\">This</a> is a great post on the structure tensor, by Chris Luengo.</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3b02791",
   "metadata": {},
   "source": [
    "The structure tensor is the outer product of the gradient vector with itself, locally averaged. \n",
    "\n",
    "$$\\mathbf{S} = \\overline{(\\nabla \\mathbf{I})(\\nabla \\mathbf{I})^{\\top}} = \\overline{\\begin{pmatrix} \\mathbf{I}_{x} \\\\\n",
    "\\mathbf{I}_{y}\\end{pmatrix}  \\begin{pmatrix} \\mathbf{I}_{x} &\n",
    "\\mathbf{I}_{y}\\end{pmatrix}} = \\begin{pmatrix}\n",
    "\\overline{\\mathbf{I}_{x}\\mathbf{I}_{x}} & \\overline{\\mathbf{I}_{x}\\mathbf{I}_{y}} \\\\\n",
    "\\overline{\\mathbf{I}_{x}\\mathbf{I}_{y}} & \\overline{\\mathbf{I}_{y}\\mathbf{I}_{y}} \n",
    "\\end{pmatrix}$$\n",
    "\n",
    "where $\\mathbf{I}_{x}$ indicates the partial derivative along axis $x$, and the overlines $\\overline{\\cdot}$ indicate local averaging, usually by means of a Gaussian kernel. \n",
    "\n",
    "*Note that the structure tensor is composed by **first-order** partial derivatives.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5243444",
   "metadata": {},
   "source": [
    "## 1) The image gradient\n",
    "To compute the structure tensor, we first calculate the gradient of the image at each point, which gives us a vector that describes the direction and strength of the local intensity change. \n",
    "\n",
    "Because an image has usually two dimensions, the gradient consists on the stacking of two partial derivatives:\n",
    "\n",
    "$$ \\nabla I(x, y) = \\begin{pmatrix} \n",
    "                            \\frac{\\partial I(x, y) }{\\partial x }\\\\\n",
    "                            \\frac{\\partial I(x, y) }{\\partial x }\n",
    "                    \\end{pmatrix} = \n",
    "                    \\begin{pmatrix} \n",
    "                            \\mathbf{I}_{x}\\\\\n",
    "                            \\mathbf{I}_{y}\n",
    "                    \\end{pmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6886771",
   "metadata": {},
   "source": [
    "Lets start with the first-order partial derivatives. \n",
    "\n",
    "The derivative of a function is defined as: \n",
    "\n",
    "$$ \\frac{\\partial f(x) }{\\partial x } =  \\lim_{h \\to 0} \\frac{f(x+h) - f(x)}{h}$$\n",
    "\n",
    "\n",
    "On a **discrete grid**, the smallest distance obtainable without interpolation is $h=1$, which yields the [Finite Difference Method](https://www.wikiwand.com/en/Finite_difference_method) approximation of the derivative:\n",
    "\n",
    "$$ \\left(\\frac{\\partial f(x) }{\\partial x }\\right)_{FD} \\approx  f(x+1) - f(x)$$\n",
    "\n",
    "\n",
    "In practice, the finite diffrence operator corresponds to a *convolution* with a linear filter with values `[1,-1]`:\n",
    "\n",
    "$$ f(x+1) - f(x) = f(x) * [1 \\, \\, -1]$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d642bea6",
   "metadata": {},
   "source": [
    "Let's use *Numba Stencils* to efficiently perform  `finite_difference`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0745b5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.stencil\n",
    "def _finite_difference_x(images):\n",
    "    \"\"\"\n",
    "    Apply finite differences on the x-axis of a 3D image.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "        images: nd-array (npix_y, npix_x, nchan)\n",
    "                Stack of images\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "        derivatives: nd-array (npix_y, npix_x, nchan)\n",
    "                Stack of images\n",
    "    \"\"\"\n",
    "    return images[0, 1, 0] - images[0, 0, 0]\n",
    "\n",
    "\n",
    "@numba.stencil\n",
    "def _finite_difference_y(images):\n",
    "    \"\"\"\n",
    "    Apply finite differences on the y-axis of a 3D image.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "        images: nd-array (npix_y, npix_x, nchan)\n",
    "                Stack of images\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "        derivatives: nd-array (npix_y, npix_x, nchan)\n",
    "                Stack of images\n",
    "    \"\"\"\n",
    "    return images[1, 0, 0] - images[0, 0, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60896cdf",
   "metadata": {},
   "source": [
    "As we saw on the first notebook, we can compile these functions to run even faster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89334e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(parallel=True)\n",
    "def fd_x(images):\n",
    "    return _finite_difference_x(images)\n",
    "\n",
    "\n",
    "@numba.jit(parallel=True)\n",
    "def fd_y(images):\n",
    "    return _finite_difference_y(images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101613d2",
   "metadata": {},
   "source": [
    "Now we can apply each function to each array chunk with [`map_overlap`](https://docs.dask.org/en/stable/generated/dask.array.map_overlap.html). Dask will run the computation in multi-threaded mode. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea305974",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_fd = hubble_image.map_overlap(fd_x, depth=(0, 1, 0))\n",
    "dfdy_fd = hubble_image.map_overlap(fd_y, depth=(1, 0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15344b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        dfdx_fd.blocks[0, 20].mean(-1),\n",
    "        dfdy_fd.blocks[0, 20].mean(-1),\n",
    "    ],\n",
    "    titles=[\"Original\", \"fd_x\", \"fd_y\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a13075b",
   "metadata": {},
   "source": [
    "The problem with the `finite_difference` filter is that it respond very strongly to the noise in the image!\n",
    "\n",
    "As a solution to that, we can use the **Gaussian derivative** (see another [great post](https://www.crisluengo.net/archives/22/) from Chris Luengo) on its benefits with respect to the finite difference method. \n",
    "\n",
    "The Gaussian derivative is defined as:\n",
    "$$\\begin{aligned}\n",
    " \\frac{\\partial}{\\partial x} [I(x, y) \\ast g(x, y)] &=  \\frac{\\partial}{\\partial x} \\ast I(x, y) \\ast g(x, y) \\\\\n",
    "                                                    &=  I(x, y) \\ast \\frac{\\partial}{\\partial x} \\ast g(x, y) \\\\\n",
    "                                                    &=  I(x, y) \\ast \\left[\\frac{\\partial}{\\partial x} g(x, y)\\right]\n",
    "\\end{aligned}$$\n",
    "\n",
    "Where:\n",
    "- in the first step we used the associative property of the convolution. \n",
    "- in the second step we used the commutative property  property of the convolution. \n",
    "- in the third step we used the associative property of the convolution. \n",
    "\n",
    "These properties show us that computing the gradient of an image blurred with a Gaussian is the same thing as convolving the image with the gradient of a Gaussian!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3d392d5",
   "metadata": {},
   "source": [
    "Let's use again *Numba stencils* to perform the `gaussian_derivative`. This time, we'll define the convolutional kernel using a Scipy function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "363322ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage._filters import _gaussian_kernel1d\n",
    "\n",
    "gaussian_derivative = _gaussian_kernel1d(1.5, 1, 10)\n",
    "plt.plot(gaussian_derivative, label=\"Gaussian Derivative kernel\")\n",
    "plt.xlabel(r\"$x$\")\n",
    "plt.ylabel(r\"$g'(x)$\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee830e8d",
   "metadata": {},
   "source": [
    "There are pre-defined some helper function in `utils` to create stencils from a 1d array kernel. Check them out if you have time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06d32fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import apply_kernel_x, apply_kernel_y\n",
    "\n",
    "\n",
    "@numba.jit(nopython=True, parallel=True)\n",
    "def gd_x(images):\n",
    "    return apply_kernel_x(images, gaussian_derivative)\n",
    "\n",
    "\n",
    "@numba.jit(nopython=True, parallel=True)\n",
    "def gd_y(images):\n",
    "    return apply_kernel_y(images, gaussian_derivative)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ca35bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_gd = hubble_image.map_overlap(gd_x, depth=(0, 10, 0))\n",
    "dfdy_gd = hubble_image.map_overlap(gd_y, depth=(10, 0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eb89acf",
   "metadata": {},
   "source": [
    "#### Can you see any difference between finite differences and Gaussian derivative?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a2f7859",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        dfdx_fd.blocks[0, 20].mean(-1),\n",
    "        dfdy_fd.blocks[0, 20].mean(-1),\n",
    "    ],\n",
    "    titles=[\"Original\", \"Finite Diff. (x-axis)\", \"Finite Diff. (y-axis)\"],\n",
    ")\n",
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        dfdx_gd.blocks[0, 20].mean(-1),\n",
    "        dfdy_gd.blocks[0, 20].mean(-1),\n",
    "    ],\n",
    "    titles=[\"Original\", \"Gauss. Der. (x-axis)\", \"Gauss. Der. (y-axis)\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f66acd0",
   "metadata": {},
   "source": [
    "## 2) The Gradient outer product\n",
    "\n",
    "$$\\mathbf{S} = \\overline{(\\nabla \\mathbf{I})(\\nabla \\mathbf{I})^{\\top}} = \\overline{\\begin{pmatrix} \\mathbf{I}_{x} \\\\\n",
    "\\mathbf{I}_{y}\\end{pmatrix}  \\begin{pmatrix} \\mathbf{I}_{x} &\n",
    "\\mathbf{I}_{y}\\end{pmatrix}} = \\begin{pmatrix}\n",
    "\\overline{\\mathbf{I}_{x}\\mathbf{I}_{x}} & \\overline{\\mathbf{I}_{x}\\mathbf{I}_{y}} \\\\\n",
    "\\overline{\\mathbf{I}_{x}\\mathbf{I}_{y}} & \\overline{\\mathbf{I}_{y}\\mathbf{I}_{y}} \n",
    "\\end{pmatrix}$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31479c14",
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = {0: \"gaussian_derivative\", 1: \"finite_differences\"}\n",
    "\n",
    "method = methods[0]\n",
    "\n",
    "if method == \"finite_differences\":\n",
    "    IxIx = dfdx_fd * dfdx_fd\n",
    "    IxIy = dfdx_fd * dfdy_fd\n",
    "    IyIy = dfdy_fd * dfdy_fd\n",
    "elif method == \"gaussian_derivative\":\n",
    "    IxIx = dfdx_gd * dfdx_gd\n",
    "    IxIy = dfdx_gd * dfdy_gd\n",
    "    IyIy = dfdy_gd * dfdy_gd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d7abf49",
   "metadata": {},
   "source": [
    "### 3) Local averaging with a Gaussian window"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6846bf3",
   "metadata": {},
   "source": [
    "We then take these gradient vectors and compute a second-order tensor that describes the covariance of the gradient vectors in a small region around each point. This is in practice performed by a weighted average using a Gaussian window like the following one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a8354e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "gaussian_kernel1d = _gaussian_kernel1d(2, 0, 10)\n",
    "plt.plot(gaussian_kernel1d, label=\"Gaussian kernel\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b3177e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.jit(nopython=True, parallel=True)\n",
    "def smooth(images):\n",
    "    return apply_kernel_y(apply_kernel_x(images, gaussian_kernel1d), gaussian_kernel1d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1fa6143",
   "metadata": {},
   "outputs": [],
   "source": [
    "IxIx_bar = IxIx.map_overlap(smooth, depth=(10, 10, 0))\n",
    "IxIy_bar = IxIy.map_overlap(smooth, depth=(10, 10, 0))\n",
    "IyIy_bar = IyIy.map_overlap(smooth, depth=(10, 10, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f96b372c",
   "metadata": {},
   "source": [
    "This tensor tells us how the gradient vectors are aligned and how strong they are in different directions, which gives us information about the local texture and patterns in the image. We can use this information to identify features like edges, corners, and lines, and to track the movement of objects in a sequence of images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16d5f161",
   "metadata": {},
   "outputs": [],
   "source": [
    "IxIx_bar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69b69642",
   "metadata": {},
   "source": [
    "### 4) The structure tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c86bb39",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        IxIx_bar.blocks[0, 20].mean(-1),\n",
    "        IxIy_bar.blocks[0, 20].mean(-1),\n",
    "        IyIy_bar.blocks[0, 20].mean(-1),\n",
    "    ],\n",
    "    titles=[\"Original\", \"IxIx_bar\", \"IxIy_bar\", \"IyIy_bar\"],\n",
    "    cmap=[None, \"pink_r\", \"pink_r\", \"pink_r\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9181b8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_tensor = da.stack(\n",
    "    [da.stack([IxIx_bar, IxIy_bar], axis=3), da.stack([IxIy_bar, IyIy_bar], axis=3)],\n",
    "    axis=4,\n",
    ")\n",
    "structure_tensor = structure_tensor.rechunk(hubble_image.chunksize + (2, 2))\n",
    "structure_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e55d95",
   "metadata": {},
   "source": [
    "Rechunking across axes can be expensive and incur a lot of communication, but Dask array has fairly efficient algorithms to accomplish this.\n",
    "\n",
    "In our case, rechunking is necessary for the next steps of the pipeline:\n",
    "\n",
    "## 5) Eigendecomposition of the structure tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fcea724",
   "metadata": {},
   "outputs": [],
   "source": [
    "w, v = da.apply_gufunc(np.linalg.eigh, \"(m,m)->(m),(m,m)\", structure_tensor)\n",
    "eigvals, eivgecs = w.persist(), v.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5732a36d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bonus exercise to deepen your Dask+Numba skills:\n",
    "# 1) Create a generalized universal function @numba.guvectorize\n",
    "### This should perform the eigendecomposition in parallel for some axis (use @numba.prange)\n",
    "# 2) Apply it to the structure tensor via dask.array.apply_gufunc.\n",
    "\n",
    "# Is it faster than the da.apply_gufunc(np.linalg.eigh) call?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101af3fc",
   "metadata": {},
   "source": [
    "## 6) Results and interpretation\n",
    "\n",
    "##### 6.1. Local gradient strength and Local gradient variation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc6c58f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[0, 20],\n",
    "        eigvals.blocks[0, 20][..., 0].mean(-1),\n",
    "        eigvals.blocks[0, 20][..., 1].mean(-1),\n",
    "    ],\n",
    "    titles=[\"Original\", \"Gradient stength\", \"Gradient Variation\"],\n",
    "    cmap=[None, \"pink_r\", \"pink_r\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7700153",
   "metadata": {},
   "source": [
    "##### 6.2. Energy and anisotropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7da7a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy = eigvals[..., 0] + eigvals[..., 1]\n",
    "anisotropy1 = (eigvals[..., 1] - eigvals[..., 0]) / energy\n",
    "anisotropy2 = 1 - (eigvals[..., 0] / eigvals[..., 1])\n",
    "\n",
    "show_images(python plus \n",
    "    images=[energy.blocks[0, 20].mean(-1), anisotropy1.blocks[0, 20].mean(-1)],\n",
    "    titles=[\"energy\", \"anisotropy1\"],\n",
    "    cmap=[\"pink_r\", \"pink_r\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a056abe6",
   "metadata": {},
   "source": [
    "##### 6.3 Feature extraction\n",
    "We can now select the features that best captue young galaxies, for example, low anisotropy and high energy in red channel would be good attributes to track."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7748ef07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# low anisotropy\n",
    "feature1 = 1 - anisotropy1[...].mean(-1)\n",
    "\n",
    "# high red, low blue and green\n",
    "feature2 = energy[..., 0] ** 2 / (energy[..., 0] + energy[..., 1] + energy[..., 2])\n",
    "\n",
    "features = feature1 * feature2\n",
    "\n",
    "show_images(\n",
    "    images=[hubble_image, features],\n",
    "    zoom=[True, True],\n",
    "    titles=[\"Original\", \"low anisotropy + high energy\"],\n",
    "    cmap=[None, \"pink_r\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db3f1cd",
   "metadata": {},
   "source": [
    "We can now filter and select the chunks that only contain high expression of those features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628c511b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "top1 = 0.008\n",
    "blocks_of_interest = []\n",
    "for bl in itertools.product(np.arange(24), np.arange(27)):\n",
    "    if da.nansum(features.blocks[bl] > top1):\n",
    "        blocks_of_interest.append(bl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03fd3f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(blocks_of_interest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c683d89",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "show_images(\n",
    "    images=[\n",
    "        hubble_image.blocks[blocks_of_interest[0]],\n",
    "        hubble_image.blocks[blocks_of_interest[1]],\n",
    "        hubble_image.blocks[blocks_of_interest[2]],\n",
    "        hubble_image.blocks[blocks_of_interest[3]],\n",
    "    ]\n",
    ")\n",
    "\n",
    "show_images(\n",
    "    images=[\n",
    "        features.blocks[blocks_of_interest[0]],\n",
    "        features.blocks[blocks_of_interest[1]],\n",
    "        features.blocks[blocks_of_interest[2]],\n",
    "        features.blocks[blocks_of_interest[3]],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e00f1b19",
   "metadata": {},
   "source": [
    "The whole picture can be seen at https://esahubble.org/images/heic0406a/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5ce36eb",
   "metadata": {},
   "source": [
    "# BONUS\n",
    "##### Create a single function that performs the ST using Dask and Numba without storing any intermediate steps (such as partial derivatives)."
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}