examples/genotypes_count_alleles.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.1.0a5.dev0+dirty'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import skallel_tensor\n",
    "from skallel_tensor import genotypes_count_alleles\n",
    "from numba import cuda\n",
    "skallel_tensor.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Smaller dataset (400MB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate some genotype data\n",
    "gt = np.random.choice(np.array([-1, 0, 1, 2, 3], dtype='i1'), \n",
    "                      p=[.01, .7, .2, .08, .01], \n",
    "                      size=(200_000, 1_000, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "400.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt.nbytes / 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('int8')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[0, 0],\n",
       "        [1, 0],\n",
       "        [0, 0],\n",
       "        ...,\n",
       "        [0, 0],\n",
       "        [2, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 2],\n",
       "        [0, 1],\n",
       "        [0, 1],\n",
       "        ...,\n",
       "        [1, 0],\n",
       "        [2, 0],\n",
       "        [2, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ...,\n",
       "        [2, 0],\n",
       "        [0, 0],\n",
       "        [1, 0]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [1, 0],\n",
       "        ...,\n",
       "        [0, 0],\n",
       "        [1, 1],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 1],\n",
       "        [1, 0],\n",
       "        ...,\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 2]],\n",
       "\n",
       "       [[0, 1],\n",
       "        [0, 1],\n",
       "        [1, 0],\n",
       "        ...,\n",
       "        [0, 0],\n",
       "        [1, 1],\n",
       "        [0, 0]]], dtype=int8)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200000, 1000, 2)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single CPU core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1394,  426,  142,   18],\n",
       "       [1424,  390,  148,   23],\n",
       "       [1444,  383,  135,   21],\n",
       "       ...,\n",
       "       [1377,  426,  166,   17],\n",
       "       [1429,  383,  149,   23],\n",
       "       [1414,  388,  145,   31]], dtype=int32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# warm-up jit\n",
    "genotypes_count_alleles(gt, max_allele=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "456 ms ± 3.16 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r3\n",
    "genotypes_count_alleles(gt, max_allele=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple CPU cores via Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.array as da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>  <thead>    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 400.00 MB </td> <td> 20.00 MB </td></tr>\n",
       "    <tr><th> Shape </th><td> (200000, 1000, 2) </td> <td> (10000, 1000, 2) </td></tr>\n",
       "    <tr><th> Count </th><td> 21 Tasks </td><td> 20 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> int8 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody></table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"156\" height=\"146\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"10\" y1=\"25\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"13\" y1=\"3\" x2=\"13\" y2=\"28\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"17\" y2=\"32\" />\n",
       "  <line x1=\"20\" y1=\"10\" x2=\"20\" y2=\"36\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"24\" y2=\"39\" />\n",
       "  <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"43\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"31\" y2=\"46\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"34\" y2=\"50\" />\n",
       "  <line x1=\"38\" y1=\"28\" x2=\"38\" y2=\"53\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"41\" y2=\"57\" />\n",
       "  <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"60\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"48\" y2=\"64\" />\n",
       "  <line x1=\"52\" y1=\"42\" x2=\"52\" y2=\"67\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"55\" y2=\"71\" />\n",
       "  <line x1=\"59\" y1=\"49\" x2=\"59\" y2=\"74\" />\n",
       "  <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"78\" />\n",
       "  <line x1=\"66\" y1=\"56\" x2=\"66\" y2=\"81\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"70\" y2=\"85\" />\n",
       "  <line x1=\"73\" y1=\"63\" x2=\"73\" y2=\"88\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"77\" y2=\"92\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 80.588235,70.588235 80.588235,96.000852 10.000000,25.412617\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"35\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"13\" y1=\"3\" x2=\"38\" y2=\"3\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"42\" y2=\"7\" />\n",
       "  <line x1=\"20\" y1=\"10\" x2=\"46\" y2=\"10\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"49\" y2=\"14\" />\n",
       "  <line x1=\"27\" y1=\"17\" x2=\"53\" y2=\"17\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"56\" y2=\"21\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"60\" y2=\"24\" />\n",
       "  <line x1=\"38\" y1=\"28\" x2=\"63\" y2=\"28\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"67\" y2=\"31\" />\n",
       "  <line x1=\"45\" y1=\"35\" x2=\"70\" y2=\"35\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"74\" y2=\"38\" />\n",
       "  <line x1=\"52\" y1=\"42\" x2=\"77\" y2=\"42\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"81\" y2=\"45\" />\n",
       "  <line x1=\"59\" y1=\"49\" x2=\"84\" y2=\"49\" />\n",
       "  <line x1=\"62\" y1=\"52\" x2=\"88\" y2=\"52\" />\n",
       "  <line x1=\"66\" y1=\"56\" x2=\"91\" y2=\"56\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"95\" y2=\"60\" />\n",
       "  <line x1=\"73\" y1=\"63\" x2=\"98\" y2=\"63\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"102\" y2=\"67\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"35\" y1=\"0\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 35.412617,0.000000 106.000852,70.588235 80.588235,70.588235\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"96\" x2=\"106\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"106\" y1=\"70\" x2=\"106\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"80.588235,70.588235 106.000852,70.588235 106.000852,96.000852 80.588235,96.000852\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"93.294544\" y=\"116.000852\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >2</text>\n",
       "  <text x=\"126.000852\" y=\"83.294544\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,126.000852,83.294544)\">1000</text>\n",
       "  <text x=\"35.294118\" y=\"80.706734\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,80.706734)\">200000</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<array, shape=(200000, 1000, 2), dtype=int8, chunksize=(10000, 1000, 2)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt_dask = da.from_array(gt, chunks=(10_000, None, None))\n",
    "gt_dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "152 ms ± 57.8 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r5\n",
    "genotypes_count_alleles(gt_dask, max_allele=3).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mon Jul 22 13:02:54 2019       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 430.26       Driver Version: 430.26       CUDA Version: 10.2     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  Quadro M1000M       Off  | 00000000:01:00.0 Off |                  N/A |\r\n",
      "| N/A   49C    P0    N/A /  N/A |    933MiB /  2004MiB |      5%      Default |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi | head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 70 ms, sys: 104 ms, total: 174 ms\n",
      "Wall time: 173 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<numba.cuda.cudadrv.devicearray.DeviceNDArray at 0x7f079bb01b70>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "gt_cuda = cuda.to_device(gt)\n",
    "gt_cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1394,  426,  142,   18],\n",
       "       [1424,  390,  148,   23],\n",
       "       [1444,  383,  135,   21],\n",
       "       ...,\n",
       "       [1377,  426,  166,   17],\n",
       "       [1429,  383,  149,   23],\n",
       "       [1414,  388,  145,   31]], dtype=int32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# warm up jit\n",
    "genotypes_count_alleles(gt_cuda, max_allele=3).copy_to_host()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "181 ms ± 9.86 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r5\n",
    "genotypes_count_alleles(gt_cuda, max_allele=3)\n",
    "cuda.synchronize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<numba.cuda.cudadrv.devicearray.DeviceNDArray at 0x7f079b6d3240>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ac_cuda = genotypes_count_alleles(gt_cuda, max_allele=3)\n",
    "ac_cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 149 ms, sys: 40.2 ms, total: 189 ms\n",
      "Wall time: 189 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1394,  426,  142,   18],\n",
       "       [1424,  390,  148,   23],\n",
       "       [1444,  383,  135,   21],\n",
       "       ...,\n",
       "       [1377,  426,  166,   17],\n",
       "       [1429,  383,  149,   23],\n",
       "       [1414,  388,  145,   31]], dtype=int32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time ac_cuda.copy_to_host()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPU via Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_dask_cuda = gt_dask.map_blocks(cuda.to_device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1394,  426,  142,   18],\n",
       "       [1424,  390,  148,   23],\n",
       "       [1444,  383,  135,   21],\n",
       "       ...,\n",
       "       [1377,  426,  166,   17],\n",
       "       [1429,  383,  149,   23],\n",
       "       [1414,  388,  145,   31]], dtype=int32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_count_alleles(gt_dask_cuda, max_allele=3).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "403 ms ± 15.8 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n1 -r3\n",
    "genotypes_count_alleles(gt_dask_cuda, max_allele=3).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Larger dataset (4GB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gt_random_big = da.random.choice(\n",
    "#     np.array([-1, 0, 1, 2, 3], dtype='i1'), \n",
    "#     p=[.01, .7, .2, .08, .01], \n",
    "#     size=(2_000_000, 1_000, 2),\n",
    "#     chunks=(50_000, None, None))\n",
    "# gt_random_big"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gt_random_big.to_zarr('example.zarr', component='gt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Name</th><td style=\"text-align: left\">/gt</td></tr><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">int8</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(2000000, 1000, 2)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(50000, 1000, 2)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.DirectoryStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">4000000000 (3.7G)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">2297009415 (2.1G)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.7</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">40/40</td></tr></tbody></table>"
      ],
      "text/plain": [
       "Name               : /gt\n",
       "Type               : zarr.core.Array\n",
       "Data type          : int8\n",
       "Shape              : (2000000, 1000, 2)\n",
       "Chunk shape        : (50000, 1000, 2)\n",
       "Order              : C\n",
       "Read-only          : False\n",
       "Compressor         : Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)\n",
       "Store type         : zarr.storage.DirectoryStore\n",
       "No. bytes          : 4000000000 (3.7G)\n",
       "No. bytes stored   : 2297009415 (2.1G)\n",
       "Storage ratio      : 1.7\n",
       "Chunks initialized : 40/40"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import zarr\n",
    "gt_big_zarr = zarr.open('example.zarr')['gt']\n",
    "gt_big_zarr.info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>  <thead>    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 4.00 GB </td> <td> 100.00 MB </td></tr>\n",
       "    <tr><th> Shape </th><td> (2000000, 1000, 2) </td> <td> (50000, 1000, 2) </td></tr>\n",
       "    <tr><th> Count </th><td> 41 Tasks </td><td> 40 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> int8 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody></table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"156\" height=\"146\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"10\" y1=\"25\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"11\" y1=\"1\" x2=\"11\" y2=\"27\" />\n",
       "  <line x1=\"13\" y1=\"3\" x2=\"13\" y2=\"28\" />\n",
       "  <line x1=\"15\" y1=\"5\" x2=\"15\" y2=\"30\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"17\" y2=\"32\" />\n",
       "  <line x1=\"18\" y1=\"8\" x2=\"18\" y2=\"34\" />\n",
       "  <line x1=\"20\" y1=\"10\" x2=\"20\" y2=\"36\" />\n",
       "  <line x1=\"22\" y1=\"12\" x2=\"22\" y2=\"37\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"24\" y2=\"39\" />\n",
       "  <line x1=\"25\" y1=\"15\" x2=\"25\" y2=\"41\" />\n",
       "  <line x1=\"27\" y1=\"17\" x2=\"27\" y2=\"43\" />\n",
       "  <line x1=\"29\" y1=\"19\" x2=\"29\" y2=\"44\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"31\" y2=\"46\" />\n",
       "  <line x1=\"32\" y1=\"22\" x2=\"32\" y2=\"48\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"34\" y2=\"50\" />\n",
       "  <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"51\" />\n",
       "  <line x1=\"38\" y1=\"28\" x2=\"38\" y2=\"53\" />\n",
       "  <line x1=\"40\" y1=\"30\" x2=\"40\" y2=\"55\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"41\" y2=\"57\" />\n",
       "  <line x1=\"43\" y1=\"33\" x2=\"43\" y2=\"58\" />\n",
       "  <line x1=\"45\" y1=\"35\" x2=\"45\" y2=\"60\" />\n",
       "  <line x1=\"47\" y1=\"37\" x2=\"47\" y2=\"62\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"48\" y2=\"64\" />\n",
       "  <line x1=\"50\" y1=\"40\" x2=\"50\" y2=\"66\" />\n",
       "  <line x1=\"52\" y1=\"42\" x2=\"52\" y2=\"67\" />\n",
       "  <line x1=\"54\" y1=\"44\" x2=\"54\" y2=\"69\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"55\" y2=\"71\" />\n",
       "  <line x1=\"57\" y1=\"47\" x2=\"57\" y2=\"73\" />\n",
       "  <line x1=\"59\" y1=\"49\" x2=\"59\" y2=\"74\" />\n",
       "  <line x1=\"61\" y1=\"51\" x2=\"61\" y2=\"76\" />\n",
       "  <line x1=\"62\" y1=\"52\" x2=\"62\" y2=\"78\" />\n",
       "  <line x1=\"64\" y1=\"54\" x2=\"64\" y2=\"80\" />\n",
       "  <line x1=\"66\" y1=\"56\" x2=\"66\" y2=\"81\" />\n",
       "  <line x1=\"68\" y1=\"58\" x2=\"68\" y2=\"83\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"70\" y2=\"85\" />\n",
       "  <line x1=\"71\" y1=\"61\" x2=\"71\" y2=\"87\" />\n",
       "  <line x1=\"73\" y1=\"63\" x2=\"73\" y2=\"88\" />\n",
       "  <line x1=\"75\" y1=\"65\" x2=\"75\" y2=\"90\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"77\" y2=\"92\" />\n",
       "  <line x1=\"78\" y1=\"68\" x2=\"78\" y2=\"94\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 80.588235,70.588235 80.588235,96.000852 10.000000,25.412617\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"35\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"11\" y1=\"1\" x2=\"37\" y2=\"1\" />\n",
       "  <line x1=\"13\" y1=\"3\" x2=\"38\" y2=\"3\" />\n",
       "  <line x1=\"15\" y1=\"5\" x2=\"40\" y2=\"5\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"42\" y2=\"7\" />\n",
       "  <line x1=\"18\" y1=\"8\" x2=\"44\" y2=\"8\" />\n",
       "  <line x1=\"20\" y1=\"10\" x2=\"46\" y2=\"10\" />\n",
       "  <line x1=\"22\" y1=\"12\" x2=\"47\" y2=\"12\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"49\" y2=\"14\" />\n",
       "  <line x1=\"25\" y1=\"15\" x2=\"51\" y2=\"15\" />\n",
       "  <line x1=\"27\" y1=\"17\" x2=\"53\" y2=\"17\" />\n",
       "  <line x1=\"29\" y1=\"19\" x2=\"54\" y2=\"19\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"56\" y2=\"21\" />\n",
       "  <line x1=\"32\" y1=\"22\" x2=\"58\" y2=\"22\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"60\" y2=\"24\" />\n",
       "  <line x1=\"36\" y1=\"26\" x2=\"61\" y2=\"26\" />\n",
       "  <line x1=\"38\" y1=\"28\" x2=\"63\" y2=\"28\" />\n",
       "  <line x1=\"40\" y1=\"30\" x2=\"65\" y2=\"30\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"67\" y2=\"31\" />\n",
       "  <line x1=\"43\" y1=\"33\" x2=\"68\" y2=\"33\" />\n",
       "  <line x1=\"45\" y1=\"35\" x2=\"70\" y2=\"35\" />\n",
       "  <line x1=\"47\" y1=\"37\" x2=\"72\" y2=\"37\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"74\" y2=\"38\" />\n",
       "  <line x1=\"50\" y1=\"40\" x2=\"76\" y2=\"40\" />\n",
       "  <line x1=\"52\" y1=\"42\" x2=\"77\" y2=\"42\" />\n",
       "  <line x1=\"54\" y1=\"44\" x2=\"79\" y2=\"44\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"81\" y2=\"45\" />\n",
       "  <line x1=\"57\" y1=\"47\" x2=\"83\" y2=\"47\" />\n",
       "  <line x1=\"59\" y1=\"49\" x2=\"84\" y2=\"49\" />\n",
       "  <line x1=\"61\" y1=\"51\" x2=\"86\" y2=\"51\" />\n",
       "  <line x1=\"62\" y1=\"52\" x2=\"88\" y2=\"52\" />\n",
       "  <line x1=\"64\" y1=\"54\" x2=\"90\" y2=\"54\" />\n",
       "  <line x1=\"66\" y1=\"56\" x2=\"91\" y2=\"56\" />\n",
       "  <line x1=\"68\" y1=\"58\" x2=\"93\" y2=\"58\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"95\" y2=\"60\" />\n",
       "  <line x1=\"71\" y1=\"61\" x2=\"97\" y2=\"61\" />\n",
       "  <line x1=\"73\" y1=\"63\" x2=\"98\" y2=\"63\" />\n",
       "  <line x1=\"75\" y1=\"65\" x2=\"100\" y2=\"65\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"102\" y2=\"67\" />\n",
       "  <line x1=\"78\" y1=\"68\" x2=\"104\" y2=\"68\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"35\" y1=\"0\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 35.412617,0.000000 106.000852,70.588235 80.588235,70.588235\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"106\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"96\" x2=\"106\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"106\" y1=\"70\" x2=\"106\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"80.588235,70.588235 106.000852,70.588235 106.000852,96.000852 80.588235,96.000852\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"93.294544\" y=\"116.000852\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >2</text>\n",
       "  <text x=\"126.000852\" y=\"83.294544\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,126.000852,83.294544)\">1000</text>\n",
       "  <text x=\"35.294118\" y=\"80.706734\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,80.706734)\">2000000</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<array, shape=(2000000, 1000, 2), dtype=int8, chunksize=(50000, 1000, 2)>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt_big_dask = da.from_array(gt_big_zarr)\n",
    "gt_big_dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple CPUs via Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13.4 s, sys: 3.12 s, total: 16.5 s\n",
      "Wall time: 2.33 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1420,  390,  150,   18],\n",
       "       [1367,  432,  166,   17],\n",
       "       [1370,  423,  168,   20],\n",
       "       ...,\n",
       "       [1366,  426,  173,   19],\n",
       "       [1376,  423,  173,   15],\n",
       "       [1392,  411,  156,   23]], dtype=int32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "genotypes_count_alleles(gt_big_dask, max_allele=3).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single GPU via Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_big_dask_cuda = gt_big_dask.map_blocks(cuda.to_device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.09 s, sys: 2.9 s, total: 9 s\n",
      "Wall time: 9 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1420,  390,  150,   18],\n",
       "       [1367,  432,  166,   17],\n",
       "       [1370,  423,  168,   20],\n",
       "       ...,\n",
       "       [1366,  426,  173,   19],\n",
       "       [1376,  423,  173,   15],\n",
       "       [1392,  411,  156,   23]], dtype=int32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "genotypes_count_alleles(gt_big_dask_cuda, max_allele=3).compute(num_workers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.15 s, sys: 3.23 s, total: 9.38 s\n",
      "Wall time: 5.21 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1420,  390,  150,   18],\n",
       "       [1367,  432,  166,   17],\n",
       "       [1370,  423,  168,   20],\n",
       "       ...,\n",
       "       [1366,  426,  173,   19],\n",
       "       [1376,  423,  173,   15],\n",
       "       [1392,  411,  156,   23]], dtype=int32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "genotypes_count_alleles(gt_big_dask_cuda, max_allele=3).compute(num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.55 s, sys: 3.32 s, total: 9.87 s\n",
      "Wall time: 3.79 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1420,  390,  150,   18],\n",
       "       [1367,  432,  166,   17],\n",
       "       [1370,  423,  168,   20],\n",
       "       ...,\n",
       "       [1366,  426,  173,   19],\n",
       "       [1376,  423,  173,   15],\n",
       "       [1392,  411,  156,   23]], dtype=int32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "genotypes_count_alleles(gt_big_dask_cuda, max_allele=3).compute(num_workers=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.68 s, sys: 3.45 s, total: 10.1 s\n",
      "Wall time: 3.93 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[1420,  390,  150,   18],\n",
       "       [1367,  432,  166,   17],\n",
       "       [1370,  423,  168,   20],\n",
       "       ...,\n",
       "       [1366,  426,  173,   19],\n",
       "       [1376,  423,  173,   15],\n",
       "       [1392,  411,  156,   23]], dtype=int32)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "genotypes_count_alleles(gt_big_dask_cuda, max_allele=3).compute(num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interesting thing, even with just a single GPU, we can improve performance and get better saturation by having multiple Dask workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}