Effective_Computation_in_Physics/ch11-data-structures.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:8bbfa852d0138905f9bc791947ccb830a7e1c00d0e0469948254bb922aeb88da"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "hash('') == hash(0) == hash(0.0) == hash(False) == 0"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 1,
       "text": [
        "True"
       ]
      }
     ],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import pandas as pd"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "pd.Series([42, 43, 44], dtype='f8')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 3,
       "text": [
        "0    42\n",
        "1    43\n",
        "2    44\n",
        "dtype: float64"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "s = pd.Series([42, 43, 44], \n",
      "      index=[\"electron\", \n",
      "             \"proton\", \n",
      "             \"neutron\"])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "s"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 5,
       "text": [
        "electron    42\n",
        "proton      43\n",
        "neutron     44\n",
        "dtype: int64"
       ]
      }
     ],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "s['electron']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 6,
       "text": [
        "42"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# inclusive bounds\n",
      "s['electron':'proton']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 7,
       "text": [
        "electron    42\n",
        "proton      43\n",
        "dtype: int64"
       ]
      }
     ],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# integer indexing still OK\n",
      "s[1:]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 8,
       "text": [
        "proton     43\n",
        "neutron    44\n",
        "dtype: int64"
       ]
      }
     ],
     "prompt_number": 8
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "t = pd.Series({'electron': 6, \n",
      "              'neutron': 28, \n",
      "              'proton': 496, \n",
      "              'neutrino': 8128})"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 9
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "t"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 10,
       "text": [
        "electron       6\n",
        "neutrino    8128\n",
        "neutron       28\n",
        "proton       496\n",
        "dtype: int64"
       ]
      }
     ],
     "prompt_number": 10
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "s + t"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 11,
       "text": [
        "electron     48\n",
        "neutrino    NaN\n",
        "neutron      72\n",
        "proton      539\n",
        "dtype: float64"
       ]
      }
     ],
     "prompt_number": 11
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "df = pd.DataFrame({'S': s, 'T': t})"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 12
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "df"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>NaN</td>\n",
        "      <td>8128</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>44</td>\n",
        "      <td>28</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>43</td>\n",
        "      <td>496</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 13,
       "text": [
        "           S     T\n",
        "electron  42     6\n",
        "neutrino NaN  8128\n",
        "neutron   44    28\n",
        "proton    43   496"
       ]
      }
     ],
     "prompt_number": 13
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "df[::2]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>44</td>\n",
        "      <td>28</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 14,
       "text": [
        "           S   T\n",
        "electron  42   6\n",
        "neutron   44  28"
       ]
      }
     ],
     "prompt_number": 14
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dg = df.append(pd.DataFrame({'S': [-8128]}, index=['antineutrino']))\n",
      "dg"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>NaN</td>\n",
        "      <td>8128</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>44</td>\n",
        "      <td>28</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>43</td>\n",
        "      <td>496</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>antineutrino</th>\n",
        "      <td>-8128</td>\n",
        "      <td>NaN</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 15,
       "text": [
        "                 S     T\n",
        "electron        42     6\n",
        "neutrino       NaN  8128\n",
        "neutron         44    28\n",
        "proton          43   496\n",
        "antineutrino -8128   NaN"
       ]
      }
     ],
     "prompt_number": 15
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dh = dg.drop('neutron')\n",
      "dh"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>NaN</td>\n",
        "      <td>8128</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>43</td>\n",
        "      <td>496</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>antineutrino</th>\n",
        "      <td>-8128</td>\n",
        "      <td>NaN</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 16,
       "text": [
        "                 S     T\n",
        "electron        42     6\n",
        "neutrino       NaN  8128\n",
        "proton          43   496\n",
        "antineutrino -8128   NaN"
       ]
      }
     ],
     "prompt_number": 16
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "df.T"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>electron</th>\n",
        "      <th>neutrino</th>\n",
        "      <th>neutron</th>\n",
        "      <th>proton</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>S</th>\n",
        "      <td>42</td>\n",
        "      <td>NaN</td>\n",
        "      <td>44</td>\n",
        "      <td>43</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>T</th>\n",
        "      <td>6</td>\n",
        "      <td>8128</td>\n",
        "      <td>28</td>\n",
        "      <td>496</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 17,
       "text": [
        "   electron  neutrino  neutron  proton\n",
        "S        42       NaN       44      43\n",
        "T         6      8128       28     496"
       ]
      }
     ],
     "prompt_number": 17
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "df < 42"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>False</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>False</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>False</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>False</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 18,
       "text": [
        "              S      T\n",
        "electron  False   True\n",
        "neutrino  False  False\n",
        "neutron   False   True\n",
        "proton    False  False"
       ]
      }
     ],
     "prompt_number": 18
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# accessing a single column \n",
      "# will return a series\n",
      "df['T']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 19,
       "text": [
        "electron       6\n",
        "neutrino    8128\n",
        "neutron       28\n",
        "proton       496\n",
        "Name: T, dtype: int64"
       ]
      }
     ],
     "prompt_number": 19
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# setting a name to a series\n",
      "# or expression will add a \n",
      "# column to the frame\n",
      "df['small'] = df['T'] < 100\n",
      "df"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "      <th>small</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>NaN</td>\n",
        "      <td>8128</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>44</td>\n",
        "      <td>28</td>\n",
        "      <td>True</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>43</td>\n",
        "      <td>496</td>\n",
        "      <td>False</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 20,
       "text": [
        "           S     T  small\n",
        "electron  42     6   True\n",
        "neutrino NaN  8128  False\n",
        "neutron   44    28   True\n",
        "proton    43   496  False"
       ]
      }
     ],
     "prompt_number": 20
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# deleting a column will\n",
      "# remove it from the frame\n",
      "del df['small']\n",
      "df"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "html": [
        "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>S</th>\n",
        "      <th>T</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>electron</th>\n",
        "      <td>42</td>\n",
        "      <td>6</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutrino</th>\n",
        "      <td>NaN</td>\n",
        "      <td>8128</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>neutron</th>\n",
        "      <td>44</td>\n",
        "      <td>28</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>proton</th>\n",
        "      <td>43</td>\n",
        "      <td>496</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 21,
       "text": [
        "           S     T\n",
        "electron  42     6\n",
        "neutrino NaN  8128\n",
        "neutron   44    28\n",
        "proton    43   496"
       ]
      }
     ],
     "prompt_number": 21
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from blist import sorteddict"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 22
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "b = sorteddict(first=\"Albert\", \n",
      "               last=\"Einstein\",\n",
      "               birthday=[1879, 3, 14])\n",
      "b"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 23,
       "text": [
        "sorteddict({'birthday': [1879, 3, 14], 'first': 'Albert', 'last': 'Einstein'})"
       ]
      }
     ],
     "prompt_number": 23
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "b['died'] = [1955, 4, 18]\n",
      "b"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 24,
       "text": [
        "sorteddict({'birthday': [1879, 3, 14], 'died': [1955, 4, 18], 'first': 'Albert', 'last': 'Einstein'})"
       ]
      }
     ],
     "prompt_number": 24
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "list(b.keys())"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 25,
       "text": [
        "['birthday', 'died', 'first', 'last']"
       ]
      }
     ],
     "prompt_number": 25
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class Node(object):\n",
      "    \n",
      "    def __init__(self, point, left=None, right=None):\n",
      "        self.point = point\n",
      "        self.left = left\n",
      "        self.right = right\n",
      "        \n",
      "    def __repr__(self):\n",
      "        isleaf = self.left is None and self.right is None\n",
      "        s = repr(self.point)\n",
      "        if not isleaf:\n",
      "            s = \"[\" + s + \":\"\n",
      "        if self.left is not None:\n",
      "            s += \"\\n  left = \" + \"\\n  \".join(repr(self.left).split('\\n'))\n",
      "        if self.right is not None:\n",
      "            s += \"\\n  right = \" + \"\\n  \".join(repr(self.right).split('\\n'))\n",
      "        if not isleaf:\n",
      "            s += \"\\n  ]\"\n",
      "        return s\n",
      "\n",
      "\n",
      "def kdtree(points, depth=0):\n",
      "    if len(points) == 0:\n",
      "        return None\n",
      "    k = len(points[0])\n",
      "    a = depth % k\n",
      "    points = sorted(points, key=lambda x: x[a])\n",
      "    i = int(len(points) / 2)  # middle index, rounded down\n",
      "    node_left = kdtree(points[:i], depth + 1)\n",
      "    node_right = kdtree(points[i+1:], depth + 1)\n",
      "    node = Node(points[i], node_left, node_right)\n",
      "    return node"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 26
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "points = [(1, 2), (3, 2), \n",
      "          (5, 5), (2, 1), \n",
      "          (4, 3), (1, 5)]\n",
      "root = kdtree(points)\n",
      "print(root)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "[(3, 2):\n",
        "  left = [(1, 2):\n",
        "    left = (2, 1)\n",
        "    right = (1, 5)\n",
        "    ]\n",
        "  right = [(5, 5):\n",
        "    left = (4, 3)\n",
        "    ]\n",
        "  ]\n"
       ]
      }
     ],
     "prompt_number": 27
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from scipy.spatial import KDTree\n",
      "tree = KDTree(points)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 28
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tree.data"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 29,
       "text": [
        "array([[1, 2],\n",
        "       [3, 2],\n",
        "       [5, 5],\n",
        "       [2, 1],\n",
        "       [4, 3],\n",
        "       [1, 5]])"
       ]
      }
     ],
     "prompt_number": 29
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# query() defaults to only the closest point\n",
      "dist, idx = tree.query([(4.5, 1.25)])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 30
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dist"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 31,
       "text": [
        "array([ 1.67705098])"
       ]
      }
     ],
     "prompt_number": 31
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "idx "
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 32,
       "text": [
        "array([1])"
       ]
      }
     ],
     "prompt_number": 32
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# fancy index by idx to get the point\n",
      "tree.data[idx]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 33,
       "text": [
        "array([[3, 2]])"
       ]
      }
     ],
     "prompt_number": 33
    }
   ],
   "metadata": {}
  }
 ]
}