04-cupy-and-dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "# Fun exercise to demonstrate CuPy and Dask\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def png(rgb):   # Numpy is concise enough to create PNG files from raw bytes in one screenful of code\n",
    "    def chunk(tag, data):\n",
    "        out = numpy.empty(4 + len(tag) + len(data) + 4, \"u1\")   # each chunk consists of:\n",
    "        out[0:4] = numpy.array([len(data)], \">u4\").view(\"u1\")   #   4-byte data size\n",
    "        out[4:8] = numpy.frombuffer(tag, \"u1\")                  #   4-byte tag (data type)\n",
    "        out[8:8 + len(data)] = numpy.frombuffer(data, \"u1\")     #   the data\n",
    "        crc = zlib.crc32(data, zlib.crc32(tag))                 #   CRC: cyclic redundancy check\n",
    "        out[-4:] = numpy.array([crc & 0xffffffff], \">u4\").view(\"u1\")\n",
    "        return out\n",
    "\n",
    "    preamble = numpy.frombuffer(b\"\\x89PNG\\r\\n\\x1a\\n\", \"u1\")     # preamble: \"this is a PNG file\"\n",
    "    width_height = numpy.array(rgb.shape[1::-1], \">u4\")         # header: size + image type (RGBA)\n",
    "    header = numpy.concatenate([width_height.view(\"u1\"), numpy.array([8, 6, 0, 0, 0], \"u1\")])\n",
    "    headerchunk = chunk(b\"IHDR\", header.tostring())\n",
    "\n",
    "    mini, maxi = rgb.min(), rgb.max()                           # normalize the value ranges\n",
    "    rgb = ((255 / (maxi - mini))*(rgb - mini)).astype(\"u1\")     # and convert to one byte per channel\n",
    "    data = numpy.empty((rgb.shape[0], 1 + 4*rgb.shape[1]), \"u1\")\n",
    "    data[:, 0], data[:, 4::4] = 0, 255                          # beginning of scanline and opacity\n",
    "    data[:, 1::4], data[:, 2::4], data[:, 3::4] = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]  # R, G, B\n",
    "    datachunk = chunk(b\"IDAT\", zlib.compress(data.tostring()))  # scanlines are zlib-compressed\n",
    "\n",
    "    endchunk = chunk(b\"IEND\", b\"\")                              # IEND means no more chunks\n",
    "    return numpy.concatenate([preamble, headerchunk, datachunk, endchunk])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy, zlib, IPython.display\n",
    "\n",
    "def show(rgb):                  # draw the image in Jupyter\n",
    "    return IPython.display.display(IPython.display.Image(data=png(rgb)))\n",
    "\n",
    "def save(filename, rgb):        # save the image to a file\n",
    "    with open(filename, \"wb\") as file:\n",
    "        file.write(png(rgb))\n",
    "\n",
    "rgb = numpy.zeros((9, 18, 3))   # little example: what do you think it will be?\n",
    "rgb[:4, :9, 2] = rgb[4:, :, :] = rgb[:, 9:, :] = 1\n",
    "rgb[::2, :, 1:] = 0\n",
    "show(rgb); print(rgb[:, :, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb = numpy.zeros((300, 600, 3))\n",
    "y, x = numpy.mgrid[0:300, 0:600]\n",
    "rgb[:, :, 0] = numpy.exp(0 - (x - 300)**2/100**2 - (y - 100)**2/100**2)\n",
    "rgb[:, :, 1] = numpy.exp(0 - (x - 200)**2/100**2 - (y - 200)**2/100**2)\n",
    "rgb[:, :, 2] = numpy.exp(0 - (x - 400)**2/100**2 - (y - 200)**2/100**2)\n",
    "show(rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb = numpy.zeros((300, 600, 3))\n",
    "y, x = numpy.mgrid[0:300, 0:600]\n",
    "rgb[20:220, 20:520, 0] = (numpy.sin(((x +  0)/100)**2 + ((y +  0)/100)**2))[20:220, 20:520]\n",
    "rgb[50:250, 80:580, 1] = (numpy.sin(((x + 10)/100)**2 + ((y +  0)/100)**2))[50:250, 80:580]\n",
    "rgb[80:280, 50:550, 2] = (numpy.sin(((x +  0)/100)**2 + ((y + 25)/100)**2))[80:280, 50:550]\n",
    "show(rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb = numpy.zeros((300, 600, 3)); y, x = numpy.mgrid[0:300, 0:600]\n",
    "rgb[:, :, 0] = numpy.sin((x +  0)/50)**2 + numpy.cos((y +  0)/50)**2\n",
    "rgb[:, :, 1] = numpy.sin((x + 25)/50)**2 + numpy.cos((y +  0)/50)**2\n",
    "rgb[:, :, 2] = numpy.sin((x +  0)/50)**2 + numpy.cos((y + 25)/50)**2\n",
    "rgb[rgb > 1.3] = numpy.sqrt(rgb[rgb > 1.3])\n",
    "show(rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: make something pretty! Then send it to me in the Google Doc.\n",
    "\n",
    "# rgb = numpy.zeros((300, 600, 3))\n",
    "# y, x = numpy.mgrid[0:300, 0:600]\n",
    "# rgb[:, :, 0] = ???\n",
    "# rgb[:, :, 1] = ???\n",
    "# rgb[:, :, 2] = ???\n",
    "# show(rgb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "### CuPy: Numpy interface, implemented with CUDA for GPUs\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the arrays with CuPy means that all calculations will be performed on a GPU.\n",
    "import cupy\n",
    "rgb  = cupy.zeros((300, 600, 3))\n",
    "y, x = cupy.mgrid[0:300, 0:600]\n",
    "\n",
    "# CuPy uses __array_ufunc__ to override Numpy's ufuncs with CUDA kernels.\n",
    "rgb[:, :, 0] = numpy.sin((x +  0)/50)**2 + numpy.cos((y +  0)/50)**2\n",
    "rgb[:, :, 1] = numpy.sin((x + 25)/50)**2 + numpy.cos((y +  0)/50)**2\n",
    "rgb[:, :, 2] = numpy.sin((x +  0)/50)**2 + numpy.cos((y + 25)/50)**2\n",
    "rgb[rgb > 1.3] = numpy.sqrt(rgb[rgb > 1.3])\n",
    "\n",
    "# Transfer from GPU to CPU when cupy.asnumpy is called.\n",
    "show(cupy.asnumpy(rgb))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Will your creation run on my GPU? Let's try it!\n",
    "\n",
    "# rgb  = cupy.zeros((300, 600, 3))\n",
    "# y, x = cupy.mgrid[0:300, 0:600]\n",
    "\n",
    "# ???\n",
    "\n",
    "# show(cupy.asnumpy(rgb))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "### Dask: Numpy interface, implemented in parallel\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the arrays with Dask means that all calculations will be delayed.\n",
    "\n",
    "import dask.array\n",
    "\n",
    "# Original Numpy array\n",
    "y, x = numpy.mgrid[0:300, 0:600]\n",
    "\n",
    "# Split the work into 2 vertical tiles and 3 horizontal tiles.\n",
    "x = dask.array.from_array(x, chunks=(150, 200))\n",
    "y = dask.array.from_array(y, chunks=(150, 200))\n",
    "\n",
    "# Dask uses __array_ufunc__ to override Numpy's ufuncs with lazy-evaluation placeholders.\n",
    "r = numpy.sin(((x +  0)/100)**2 + ((y +  0)/100)**2)\n",
    "g = numpy.sin(((x + 10)/100)**2 + ((y +  0)/100)**2)\n",
    "b = numpy.sin(((x +  0)/100)**2 + ((y + 25)/100)**2)\n",
    "\n",
    "# Non-ufuncs have to go through special Dask functions.\n",
    "rgb = dask.array.stack([r, g, b], axis=-1)\n",
    "\n",
    "# The result is not numerical: it's a recording of all the steps needed to do the calculation.\n",
    "rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To actually calculate a Dask array, call compute.\n",
    "\n",
    "show(rgb.compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dask turns every array operation into a task and manages their dependencies.\n",
    "rgb.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "### Dask distributed\n",
    "\n",
    "The value of delayed execution is that the task graph can be submitted to a distributed computing cluster.\n",
    "\n",
    "Before evaluating the next cell, start a `dask-scheduler` and `dask-worker` in a terminal.\n",
    "\n",
    "```bash\n",
    "dask-scheduler &\n",
    "dask-worker --nthreads 8 127.0.0.1:8786 &\n",
    "```\n",
    "\n",
    "(`127.0.0.1` is this computer and `8786` is the default port number for the scheduler. Don't try this on the Binder instance.)\n",
    "\n",
    "<br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.distributed\n",
    "\n",
    "client = dask.distributed.Client(\"127.0.0.1:8786\")\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Watch the client dashboard while you run compute.\n",
    "\n",
    "show(rgb.compute())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}