notebooks/dask_2d_subset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has some profiling of Dask used to make a selection along both first and second axes of a large-ish multidimensional array. The use case is making selections of genotype data, e.g., as required for making a web-browser for genotype data as in www.malariagen.net/apps/ag1000g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zarr 2.1.1\n",
      "dask 0.11.0\n"
     ]
    }
   ],
   "source": [
    "import zarr; print('zarr', zarr.__version__)\n",
    "import dask; print('dask', dask.__version__)\n",
    "import dask.array as da\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Group(/, 8)\n",
       "  arrays: 1; samples\n",
       "  groups: 7; 2L, 2R, 3L, 3R, UNKN, X, Y_unplaced\n",
       "  store: DirectoryStore"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# here's the real data\n",
    "callset = zarr.open_group('/kwiat/2/coluzzi/ag1000g/data/phase1/release/AR3.1/variation/main/zarr2/zstd/ag1000g.phase1.ar3',\n",
    "                          mode='r')\n",
    "callset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(/3R/calldata/genotype, (22632425, 765, 2), int8, chunks=(13107, 40, 2), order=C)\n",
       "  nbytes: 32.2G; nbytes_stored: 1.0G; ratio: 31.8; initialized: 34540/34540\n",
       "  compressor: Blosc(cname='zstd', clevel=1, shuffle=2)\n",
       "  store: DirectoryStore"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# here's the array we're going to work with\n",
    "g = callset['3R/calldata/genotype']\n",
    "g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4 ms, sys: 0 ns, total: 4 ms\n",
      "Wall time: 5.13 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "dask.array<array-b..., shape=(22632425, 765, 2), dtype=int8, chunksize=(13107, 765, 2)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# wrap as dask array with very simple chunking of first dim only\n",
    "%time gd = da.from_array(g, chunks=(g.chunks[0], None, None))\n",
    "gd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((22632425,), dtype('bool'), 13167162)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load condition used to make selection on first axis\n",
    "dim0_condition = callset['3R/variants/FILTER_PASS'][:]\n",
    "dim0_condition.shape, dim0_condition.dtype, np.count_nonzero(dim0_condition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# invent a random selection for second axis\n",
    "dim1_indices = sorted(np.random.choice(765, size=100, replace=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 15.3 s, sys: 256 ms, total: 15.5 s\n",
      "Wall time: 15.5 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "dask.array<getitem..., shape=(13167162, 100, 2), dtype=int8, chunksize=(8873, 100, 2)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# setup the 2D selection - this is the slow bit\n",
    "%time gd_sel = gd[dim0_condition][:, dim1_indices]\n",
    "gd_sel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.21 s, sys: 152 ms, total: 1.36 s\n",
      "Wall time: 316 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       ..., \n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 1],\n",
       "        [0, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]],\n",
       "\n",
       "       [[0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        ..., \n",
       "        [0, 0],\n",
       "        [0, 0],\n",
       "        [0, 0]]], dtype=int8)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# now load a slice from this new selection - quick!\n",
    "%time gd_sel[1000000:1100000].compute(optimize_graph=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         105406881 function calls (79072145 primitive calls) in 26.182 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "13167268/6    6.807    0.000    9.038    1.506 slicing.py:623(check_index)\n",
      "        2    4.713    2.356    5.831    2.916 slicing.py:398(partition_by_size)\n",
      "13167270/2    4.470    0.000    8.763    4.382 slicing.py:540(posify_index)\n",
      " 52669338    4.118    0.000    4.119    0.000 {built-in method builtins.isinstance}\n",
      "        2    2.406    1.203    8.763    4.382 slicing.py:563(<listcomp>)\n",
      "        1    0.875    0.875    0.875    0.875 slicing.py:44(<listcomp>)\n",
      " 13182474    0.600    0.000    0.600    0.000 {built-in method builtins.len}\n",
      "        2    0.527    0.264    0.527    0.264 slicing.py:420(issorted)\n",
      " 13189168    0.520    0.000    0.520    0.000 {method 'append' of 'list' objects}\n",
      "        2    0.271    0.136    0.271    0.136 slicing.py:479(<listcomp>)\n",
      "        2    0.220    0.110    0.220    0.110 {built-in method builtins.sorted}\n",
      "        1    0.162    0.162    0.162    0.162 {method 'tolist' of 'numpy.ndarray' objects}\n",
      "        2    0.113    0.056   26.071   13.035 core.py:1024(__getitem__)\n",
      "        2    0.112    0.056    6.435    3.217 slicing.py:441(take_sorted)\n",
      "        1    0.111    0.111   26.182   26.182 <string>:1(<module>)\n",
      "        2    0.060    0.030   24.843   12.422 slicing.py:142(slice_with_newaxes)\n",
      "    106/3    0.039    0.000    1.077    0.359 slicing.py:15(sanitize_index)\n",
      "        3    0.037    0.012    0.037    0.012 {built-in method _hashlib.openssl_md5}\n",
      "     6726    0.012    0.000    0.017    0.000 slicing.py:567(insert_many)\n",
      "     3364    0.004    0.000    0.021    0.000 slicing.py:156(<genexpr>)\n",
      "    20178    0.003    0.000    0.003    0.000 {method 'pop' of 'list' objects}\n",
      "        8    0.000    0.000    0.000    0.000 {method 'update' of 'dict' objects}\n",
      "        2    0.000    0.000   25.920   12.960 slicing.py:60(slice_array)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:162(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:464(<listcomp>)\n",
      "    106/4    0.000    0.000    0.037    0.009 utils.py:502(__call__)\n",
      "      100    0.000    0.000    0.000    0.000 arrayprint.py:340(array2string)\n",
      "        2    0.000    0.000    0.037    0.019 base.py:343(tokenize)\n",
      "      100    0.000    0.000    0.000    0.000 {built-in method builtins.repr}\n",
      "        2    0.000    0.000   24.763   12.381 slicing.py:170(slice_wrap_lists)\n",
      "      108    0.000    0.000    0.000    0.000 abc.py:178(__instancecheck__)\n",
      "        2    0.000    0.000    6.962    3.481 slicing.py:487(take)\n",
      "        1    0.000    0.000   26.182   26.182 {built-in method builtins.exec}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:465(<listcomp>)\n",
      "        1    0.000    0.000    0.037    0.037 base.py:314(normalize_array)\n",
      "      2/1    0.000    0.000    0.000    0.000 base.py:270(normalize_seq)\n",
      "      116    0.000    0.000    0.000    0.000 _weakrefset.py:70(__contains__)\n",
      "      100    0.000    0.000    0.000    0.000 numeric.py:1835(array_str)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:47(<listcomp>)\n",
      "        6    0.000    0.000    0.000    0.000 {built-in method builtins.sum}\n",
      "        2    0.000    0.000    0.000    0.000 exceptions.py:15(merge)\n",
      "      100    0.000    0.000    0.000    0.000 inspect.py:441(getmro)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:475(<listcomp>)\n",
      "        4    0.000    0.000    0.000    0.000 dicttoolz.py:19(merge)\n",
      "        4    0.000    0.000    0.000    0.000 functoolz.py:217(__call__)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1455(normalize_chunks)\n",
      "        4    0.000    0.000    0.000    0.000 dicttoolz.py:11(_get_factory)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:467(<listcomp>)\n",
      "      100    0.000    0.000    0.000    0.000 {method 'item' of 'numpy.ndarray' objects}\n",
      "        2    0.000    0.000    0.000    0.000 core.py:794(__init__)\n",
      "        8    0.000    0.000    0.000    0.000 {built-in method builtins.all}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:197(<genexpr>)\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:183(<genexpr>)\n",
      "        5    0.000    0.000    0.000    0.000 core.py:1043(<genexpr>)\n",
      "        7    0.000    0.000    0.000    0.000 {built-in method builtins.hasattr}\n",
      "        5    0.000    0.000    0.000    0.000 slicing.py:125(<genexpr>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'view' of 'numpy.ndarray' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:192(<listcomp>)\n",
      "        3    0.000    0.000    0.000    0.000 {method 'hexdigest' of '_hashlib.HASH' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:606(replace_ellipsis)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:613(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'ravel' of 'numpy.ndarray' objects}\n",
      "        4    0.000    0.000    0.000    0.000 {method 'items' of 'dict' objects}\n",
      "        2    0.000    0.000    0.000    0.000 {method 'encode' of 'str' objects}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:207(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:826(_get_chunks)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1452(<lambda>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:149(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:150(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 functoolz.py:11(identity)\n",
      "        4    0.000    0.000    0.000    0.000 {method 'pop' of 'dict' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}\n",
      "        2    0.000    0.000    0.000    0.000 {method 'count' of 'tuple' objects}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# what's taking so long?\n",
    "import cProfile\n",
    "cProfile.run('gd[dim0_condition][:, dim1_indices]', sort='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         105406881 function calls (79072145 primitive calls) in 25.630 seconds\n",
      "\n",
      "   Ordered by: cumulative time\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "        1    0.000    0.000   25.630   25.630 {built-in method builtins.exec}\n",
      "        1    0.107    0.107   25.630   25.630 <string>:1(<module>)\n",
      "        2    0.102    0.051   25.523   12.761 core.py:1024(__getitem__)\n",
      "        2    0.001    0.000   25.381   12.691 slicing.py:60(slice_array)\n",
      "        2    0.049    0.024   24.214   12.107 slicing.py:142(slice_with_newaxes)\n",
      "        2    0.000    0.000   24.147   12.073 slicing.py:170(slice_wrap_lists)\n",
      "13167268/6    6.664    0.000    8.855    1.476 slicing.py:623(check_index)\n",
      "13167270/2    4.354    0.000    8.466    4.233 slicing.py:540(posify_index)\n",
      "        2    2.277    1.139    8.465    4.233 slicing.py:563(<listcomp>)\n",
      "        2    0.000    0.000    6.826    3.413 slicing.py:487(take)\n",
      "        2    0.111    0.056    6.331    3.165 slicing.py:441(take_sorted)\n",
      "        2    4.628    2.314    5.704    2.852 slicing.py:398(partition_by_size)\n",
      " 52669338    4.026    0.000    4.026    0.000 {built-in method builtins.isinstance}\n",
      "    106/3    0.071    0.001    1.167    0.389 slicing.py:15(sanitize_index)\n",
      "        1    0.943    0.943    0.943    0.943 slicing.py:44(<listcomp>)\n",
      " 13182474    0.581    0.000    0.581    0.000 {built-in method builtins.len}\n",
      " 13189168    0.497    0.000    0.497    0.000 {method 'append' of 'list' objects}\n",
      "        2    0.495    0.248    0.495    0.248 slicing.py:420(issorted)\n",
      "        2    0.281    0.141    0.281    0.141 slicing.py:479(<listcomp>)\n",
      "        2    0.234    0.117    0.234    0.117 {built-in method builtins.sorted}\n",
      "        1    0.152    0.152    0.152    0.152 {method 'tolist' of 'numpy.ndarray' objects}\n",
      "        2    0.000    0.000    0.039    0.020 base.py:343(tokenize)\n",
      "    106/4    0.000    0.000    0.039    0.010 utils.py:502(__call__)\n",
      "        1    0.000    0.000    0.039    0.039 base.py:314(normalize_array)\n",
      "        3    0.039    0.013    0.039    0.013 {built-in method _hashlib.openssl_md5}\n",
      "     3364    0.003    0.000    0.019    0.000 slicing.py:156(<genexpr>)\n",
      "     6726    0.012    0.000    0.016    0.000 slicing.py:567(insert_many)\n",
      "    20178    0.003    0.000    0.003    0.000 {method 'pop' of 'list' objects}\n",
      "        4    0.000    0.000    0.000    0.000 dicttoolz.py:19(merge)\n",
      "        8    0.000    0.000    0.000    0.000 {method 'update' of 'dict' objects}\n",
      "        4    0.000    0.000    0.000    0.000 functoolz.py:217(__call__)\n",
      "        2    0.000    0.000    0.000    0.000 exceptions.py:15(merge)\n",
      "      2/1    0.000    0.000    0.000    0.000 base.py:270(normalize_seq)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:162(<genexpr>)\n",
      "      100    0.000    0.000    0.000    0.000 {built-in method builtins.repr}\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:47(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:464(<listcomp>)\n",
      "      100    0.000    0.000    0.000    0.000 numeric.py:1835(array_str)\n",
      "      100    0.000    0.000    0.000    0.000 arrayprint.py:340(array2string)\n",
      "      108    0.000    0.000    0.000    0.000 abc.py:178(__instancecheck__)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:465(<listcomp>)\n",
      "        8    0.000    0.000    0.000    0.000 {built-in method builtins.all}\n",
      "        2    0.000    0.000    0.000    0.000 core.py:794(__init__)\n",
      "      116    0.000    0.000    0.000    0.000 _weakrefset.py:70(__contains__)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1455(normalize_chunks)\n",
      "        6    0.000    0.000    0.000    0.000 {built-in method builtins.sum}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:183(<genexpr>)\n",
      "      100    0.000    0.000    0.000    0.000 {method 'item' of 'numpy.ndarray' objects}\n",
      "      100    0.000    0.000    0.000    0.000 inspect.py:441(getmro)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'encode' of 'str' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:606(replace_ellipsis)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:475(<listcomp>)\n",
      "        5    0.000    0.000    0.000    0.000 slicing.py:125(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:467(<listcomp>)\n",
      "        3    0.000    0.000    0.000    0.000 {method 'hexdigest' of '_hashlib.HASH' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'view' of 'numpy.ndarray' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:192(<listcomp>)\n",
      "        4    0.000    0.000    0.000    0.000 dicttoolz.py:11(_get_factory)\n",
      "        5    0.000    0.000    0.000    0.000 core.py:1043(<genexpr>)\n",
      "        7    0.000    0.000    0.000    0.000 {built-in method builtins.hasattr}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:207(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:613(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:149(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'ravel' of 'numpy.ndarray' objects}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:197(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:826(_get_chunks)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1452(<lambda>)\n",
      "        4    0.000    0.000    0.000    0.000 {method 'pop' of 'dict' objects}\n",
      "        4    0.000    0.000    0.000    0.000 {method 'items' of 'dict' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:150(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'count' of 'tuple' objects}\n",
      "        1    0.000    0.000    0.000    0.000 functoolz.py:11(identity)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cProfile.run('gd[dim0_condition][:, dim1_indices]', sort='cumtime')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array((20000000, 200, 2), int8, chunks=(10000, 100, 2), order=C)\n",
       "  nbytes: 7.5G; nbytes_stored: 2.7G; ratio: 2.8; initialized: 4000/4000\n",
       "  compressor: Blosc(cname='zstd', clevel=1, shuffle=2)\n",
       "  store: dict"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create a synthetic dataset for profiling\n",
    "a = zarr.array(np.random.randint(-1, 4, size=(20000000, 200, 2), dtype='i1'),\n",
    "               chunks=(10000, 100, 2), compressor=zarr.Blosc(cname='zstd', clevel=1, shuffle=2))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# create a synthetic selection for first axis\n",
    "c = np.random.randint(0, 2, size=a.shape[0], dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# create a synthetic selection for second axis\n",
    "s = sorted(np.random.choice(a.shape[1], size=100, replace=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 208 ms, sys: 0 ns, total: 208 ms\n",
      "Wall time: 206 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "dask.array<array-5..., shape=(20000000, 200, 2), dtype=int8, chunksize=(10000, 200, 2)>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time d = da.from_array(a, chunks=(a.chunks[0], None, None))\n",
    "d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12 s, sys: 200 ms, total: 12.2 s\n",
      "Wall time: 12.2 s\n"
     ]
    }
   ],
   "source": [
    "%time ds = d[c][:, s]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         80095589 function calls (60091843 primitive calls) in 19.467 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "10001773/6    4.872    0.000    6.456    1.076 slicing.py:623(check_index)\n",
      "        2    3.517    1.758    4.357    2.179 slicing.py:398(partition_by_size)\n",
      "10001775/2    3.354    0.000    6.484    3.242 slicing.py:540(posify_index)\n",
      " 40007358    2.965    0.000    2.965    0.000 {built-in method builtins.isinstance}\n",
      "        2    1.749    0.875    6.484    3.242 slicing.py:563(<listcomp>)\n",
      "        1    0.878    0.878    0.878    0.878 slicing.py:44(<listcomp>)\n",
      " 10019804    0.451    0.000    0.451    0.000 {built-in method builtins.len}\n",
      " 10027774    0.392    0.000    0.392    0.000 {method 'append' of 'list' objects}\n",
      "        2    0.363    0.181    0.363    0.181 slicing.py:420(issorted)\n",
      "        2    0.270    0.135    4.786    2.393 slicing.py:441(take_sorted)\n",
      "        1    0.207    0.207    0.207    0.207 {method 'tolist' of 'numpy.ndarray' objects}\n",
      "        2    0.158    0.079    0.158    0.079 {built-in method builtins.sorted}\n",
      "        1    0.094    0.094   19.467   19.467 <string>:1(<module>)\n",
      "        2    0.079    0.040   19.373    9.686 core.py:1024(__getitem__)\n",
      "        2    0.035    0.017   18.147    9.074 slicing.py:142(slice_with_newaxes)\n",
      "        3    0.033    0.011    0.033    0.011 {built-in method _hashlib.openssl_md5}\n",
      "    106/3    0.028    0.000    1.112    0.371 slicing.py:15(sanitize_index)\n",
      "     8002    0.015    0.000    0.020    0.000 slicing.py:567(insert_many)\n",
      "     4002    0.004    0.000    0.023    0.000 slicing.py:156(<genexpr>)\n",
      "    24006    0.003    0.000    0.003    0.000 {method 'pop' of 'list' objects}\n",
      "        8    0.001    0.000    0.001    0.000 {method 'update' of 'dict' objects}\n",
      "        2    0.001    0.000    0.001    0.000 slicing.py:479(<listcomp>)\n",
      "        2    0.000    0.000   19.259    9.630 slicing.py:60(slice_array)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:162(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:464(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:465(<listcomp>)\n",
      "    106/4    0.000    0.000    0.034    0.008 utils.py:502(__call__)\n",
      "        2    0.000    0.000   18.089    9.044 slicing.py:170(slice_wrap_lists)\n",
      "      100    0.000    0.000    0.000    0.000 arrayprint.py:340(array2string)\n",
      "      100    0.000    0.000    0.000    0.000 {built-in method builtins.repr}\n",
      "      108    0.000    0.000    0.000    0.000 abc.py:178(__instancecheck__)\n",
      "        2    0.000    0.000    5.149    2.574 slicing.py:487(take)\n",
      "        2    0.000    0.000    0.034    0.017 base.py:343(tokenize)\n",
      "        1    0.000    0.000    0.033    0.033 base.py:314(normalize_array)\n",
      "      116    0.000    0.000    0.000    0.000 _weakrefset.py:70(__contains__)\n",
      "      2/1    0.000    0.000    0.000    0.000 base.py:270(normalize_seq)\n",
      "        6    0.000    0.000    0.000    0.000 {built-in method builtins.sum}\n",
      "      100    0.000    0.000    0.000    0.000 numeric.py:1835(array_str)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:47(<listcomp>)\n",
      "        1    0.000    0.000   19.467   19.467 {built-in method builtins.exec}\n",
      "      100    0.000    0.000    0.000    0.000 inspect.py:441(getmro)\n",
      "        8    0.000    0.000    0.000    0.000 {built-in method builtins.all}\n",
      "        4    0.000    0.000    0.001    0.000 dicttoolz.py:19(merge)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1455(normalize_chunks)\n",
      "      100    0.000    0.000    0.000    0.000 {method 'item' of 'numpy.ndarray' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:475(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:794(__init__)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:467(<listcomp>)\n",
      "        3    0.000    0.000    0.000    0.000 {method 'hexdigest' of '_hashlib.HASH' objects}\n",
      "        2    0.000    0.000    0.001    0.000 exceptions.py:15(merge)\n",
      "        7    0.000    0.000    0.000    0.000 {built-in method builtins.hasattr}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:606(replace_ellipsis)\n",
      "        4    0.000    0.000    0.001    0.000 functoolz.py:217(__call__)\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:183(<genexpr>)\n",
      "        4    0.000    0.000    0.000    0.000 dicttoolz.py:11(_get_factory)\n",
      "        5    0.000    0.000    0.000    0.000 core.py:1043(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'encode' of 'str' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'view' of 'numpy.ndarray' objects}\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:197(<genexpr>)\n",
      "        5    0.000    0.000    0.000    0.000 slicing.py:125(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:192(<listcomp>)\n",
      "        8    0.000    0.000    0.000    0.000 slicing.py:207(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:613(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'count' of 'tuple' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'ravel' of 'numpy.ndarray' objects}\n",
      "        1    0.000    0.000    0.000    0.000 functoolz.py:11(identity)\n",
      "        4    0.000    0.000    0.000    0.000 {method 'pop' of 'dict' objects}\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:150(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:826(_get_chunks)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1452(<lambda>)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:149(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}\n",
      "        4    0.000    0.000    0.000    0.000 {method 'items' of 'dict' objects}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cProfile.run('d[c][:, s]', sort='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 452 ms, sys: 8 ms, total: 460 ms\n",
      "Wall time: 148 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[[ 2, -1],\n",
       "        [ 2,  3],\n",
       "        [ 3,  0],\n",
       "        ..., \n",
       "        [ 1,  3],\n",
       "        [-1, -1],\n",
       "        [ 1,  1]],\n",
       "\n",
       "       [[ 1, -1],\n",
       "        [ 2,  2],\n",
       "        [-1,  2],\n",
       "        ..., \n",
       "        [ 2, -1],\n",
       "        [ 1,  3],\n",
       "        [-1, -1]],\n",
       "\n",
       "       [[ 1, -1],\n",
       "        [ 2,  0],\n",
       "        [ 0,  3],\n",
       "        ..., \n",
       "        [ 2,  2],\n",
       "        [ 3,  2],\n",
       "        [ 0,  2]],\n",
       "\n",
       "       ..., \n",
       "       [[ 1,  2],\n",
       "        [ 3, -1],\n",
       "        [ 2,  1],\n",
       "        ..., \n",
       "        [ 1,  2],\n",
       "        [ 1,  0],\n",
       "        [ 2,  0]],\n",
       "\n",
       "       [[ 1,  2],\n",
       "        [ 1,  0],\n",
       "        [ 2,  3],\n",
       "        ..., \n",
       "        [-1,  2],\n",
       "        [ 3,  3],\n",
       "        [ 1, -1]],\n",
       "\n",
       "       [[-1,  3],\n",
       "        [ 2,  2],\n",
       "        [ 1,  1],\n",
       "        ..., \n",
       "        [ 3,  3],\n",
       "        [ 0,  0],\n",
       "        [ 0,  2]]], dtype=int8)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time ds[1000000:1100000].compute(optimize_graph=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         80055494 function calls (60052157 primitive calls) in 19.425 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "10001670/3    5.032    0.000    6.671    2.224 slicing.py:623(check_index)\n",
      "        1    3.459    3.459    4.272    4.272 slicing.py:398(partition_by_size)\n",
      "10001671/1    3.287    0.000    6.378    6.378 slicing.py:540(posify_index)\n",
      " 40006704    2.999    0.000    2.999    0.000 {built-in method builtins.isinstance}\n",
      "        1    1.731    1.731    6.378    6.378 slicing.py:563(<listcomp>)\n",
      "        1    0.849    0.849    0.849    0.849 slicing.py:44(<listcomp>)\n",
      " 10011685    0.433    0.000    0.433    0.000 {built-in method builtins.len}\n",
      " 10015670    0.381    0.000    0.381    0.000 {method 'append' of 'list' objects}\n",
      "        1    0.355    0.355    0.355    0.355 slicing.py:420(issorted)\n",
      "        1    0.196    0.196    0.196    0.196 {method 'tolist' of 'numpy.ndarray' objects}\n",
      "        1    0.193    0.193    0.193    0.193 slicing.py:479(<listcomp>)\n",
      "        1    0.157    0.157    0.157    0.157 {built-in method builtins.sorted}\n",
      "        1    0.085    0.085    4.707    4.707 slicing.py:441(take_sorted)\n",
      "        1    0.085    0.085   19.425   19.425 <string>:1(<module>)\n",
      "        1    0.079    0.079   19.341   19.341 core.py:1024(__getitem__)\n",
      "        1    0.034    0.034   18.157   18.157 slicing.py:142(slice_with_newaxes)\n",
      "        2    0.033    0.017    0.033    0.017 {built-in method _hashlib.openssl_md5}\n",
      "        1    0.026    0.026    1.071    1.071 slicing.py:15(sanitize_index)\n",
      "     4001    0.007    0.000    0.009    0.000 slicing.py:567(insert_many)\n",
      "     2001    0.002    0.000    0.011    0.000 slicing.py:156(<genexpr>)\n",
      "    12003    0.001    0.000    0.001    0.000 {method 'pop' of 'list' objects}\n",
      "        4    0.000    0.000    0.000    0.000 {method 'update' of 'dict' objects}\n",
      "        1    0.000    0.000   19.228   19.228 slicing.py:60(slice_array)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:464(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:162(<genexpr>)\n",
      "        1    0.000    0.000    0.033    0.033 base.py:314(normalize_array)\n",
      "        1    0.000    0.000   18.111   18.111 slicing.py:170(slice_wrap_lists)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:465(<listcomp>)\n",
      "        1    0.000    0.000    5.062    5.062 slicing.py:487(take)\n",
      "        1    0.000    0.000    0.033    0.033 base.py:343(tokenize)\n",
      "        1    0.000    0.000   19.425   19.425 {built-in method builtins.exec}\n",
      "        2    0.000    0.000    0.000    0.000 functoolz.py:217(__call__)\n",
      "        3    0.000    0.000    0.000    0.000 {built-in method builtins.sum}\n",
      "        2    0.000    0.000    0.000    0.000 abc.py:178(__instancecheck__)\n",
      "        1    0.000    0.000    0.000    0.000 core.py:1455(normalize_chunks)\n",
      "        2    0.000    0.000    0.000    0.000 dicttoolz.py:19(merge)\n",
      "        4    0.000    0.000    0.000    0.000 _weakrefset.py:70(__contains__)\n",
      "        2    0.000    0.000    0.000    0.000 dicttoolz.py:11(_get_factory)\n",
      "        1    0.000    0.000    0.000    0.000 exceptions.py:15(merge)\n",
      "        1    0.000    0.000    0.000    0.000 core.py:794(__init__)\n",
      "        4    0.000    0.000    0.000    0.000 {built-in method builtins.all}\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:467(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'view' of 'numpy.ndarray' objects}\n",
      "        4    0.000    0.000    0.000    0.000 slicing.py:183(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'hexdigest' of '_hashlib.HASH' objects}\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:606(replace_ellipsis)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:192(<listcomp>)\n",
      "        4    0.000    0.000    0.000    0.000 slicing.py:207(<genexpr>)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:475(<listcomp>)\n",
      "        2    0.000    0.000    0.033    0.017 utils.py:502(__call__)\n",
      "        2    0.000    0.000    0.000    0.000 slicing.py:125(<genexpr>)\n",
      "        2    0.000    0.000    0.000    0.000 core.py:1043(<genexpr>)\n",
      "        4    0.000    0.000    0.000    0.000 slicing.py:197(<genexpr>)\n",
      "        1    0.000    0.000    0.000    0.000 core.py:826(_get_chunks)\n",
      "        2    0.000    0.000    0.000    0.000 {built-in method builtins.hasattr}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'ravel' of 'numpy.ndarray' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}\n",
      "        1    0.000    0.000    0.000    0.000 {method 'encode' of 'str' objects}\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:613(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 core.py:1452(<lambda>)\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:149(<listcomp>)\n",
      "        2    0.000    0.000    0.000    0.000 {method 'pop' of 'dict' objects}\n",
      "        2    0.000    0.000    0.000    0.000 {method 'items' of 'dict' objects}\n",
      "        1    0.000    0.000    0.000    0.000 slicing.py:150(<listcomp>)\n",
      "        1    0.000    0.000    0.000    0.000 {method 'count' of 'tuple' objects}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# problem is in fact just the dim0 selection\n",
    "cProfile.run('d[c]', sort='time')"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}