Timing various methods of constructing list combinations.

indices.ipynb

{
 "metadata": {
  "name": "indices"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "import itertools"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def numpy_make_grid(*args):\n",
      "    # an array of ones in the overall shape, for broadcasting\n",
      "    ones = np.ones([len(arg) for arg in args])\n",
      "    # mesh grids of the index arrays\n",
      "    midx = [(ix * ones)[None] for ix in np.ix_(*args)]\n",
      "    # make into one Nx3 array\n",
      "    idx = np.concatenate(midx).reshape((len(args), -1)).transpose()\n",
      "    return idx"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def loop_make_grid(*args):\n",
      "    # find the sizes of each dimension and the total size of the\n",
      "    # final array\n",
      "    shape = [len(arg) for arg in args]\n",
      "    size = 1\n",
      "    for sh in shape:\n",
      "        size *= sh\n",
      "    # make a list of lists to hold the indices\n",
      "    l = [1 for i in xrange(len(args))]\n",
      "    idx = [l[:] for i in xrange(size)]\n",
      "    # fill in the indices\n",
      "    rep = 1\n",
      "    for aidx, arg in enumerate(args):\n",
      "        # repeat each value in the dimension based on which\n",
      "        # dimensions we've already included\n",
      "        vals = []\n",
      "        for val in arg:\n",
      "            vals.extend([val] * rep)\n",
      "        # repeat each dimension based on which dimensions we\n",
      "        # haven't already included and actually fill in the\n",
      "        # indices\n",
      "        rest = size / (rep * len(arg))\n",
      "        for vidx, val in enumerate(vals * rest):\n",
      "            idx[vidx][aidx] = val\n",
      "        rep *= len(arg)\n",
      "    return idx"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def itertools_make_grid(*args):\n",
      "    return list(itertools.product(*args))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def numpy_indices():\n",
      "    numpy_make_grid(np.arange(10), np.arange(10), np.arange(10))\n",
      "def loop_indices():\n",
      "    loop_make_grid(range(10), range(10), range(10))\n",
      "def itertools_indices():\n",
      "    itertools_make_grid(range(10), range(10), range(10))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%timeit numpy_indices()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "10000 loops, best of 3: 65.9 \u00b5s per loop\n"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%timeit loop_indices()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "1000 loops, best of 3: 815 \u00b5s per loop\n"
       ]
      }
     ],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%timeit itertools_indices()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "10000 loops, best of 3: 49.9 \u00b5s per loop\n"
       ]
      }
     ],
     "prompt_number": 8
    }
   ],
   "metadata": {}
  }
 ]
}