Spring Layout Animation Demo Notebook

spring_layout_animate.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:ccbb3b0da78da233f01c7845161f077f37b78bf53110da8466dc1b52a9a2e731"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Spring Layout Animation\n",
      "\n",
      "This notebook was used to develop a customized spring layout.  It uses Bokeh to animate the algorithm.\n",
      "\n",
      "**You will need to run this notebook to see the animation.**\n",
      "\n",
      "## Setup\n",
      "\n",
      "Here's the conda commands to setup an environment for this notebook.\n",
      "\n",
      "```bash\n",
      "conda install ipython-notebook\n",
      "conda install numba bokeh\n",
      "```\n"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import random\n",
      "import math\n",
      "import numpy as np\n",
      "\n",
      "from IPython.html.widgets import interact\n",
      "\n",
      "from numba import jit\n",
      "from bokeh import plotting as bp\n",
      "from bokeh.models import GlyphRenderer, DataRange1d\n",
      "\n",
      "\n",
      "bp.output_notebook()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Populate Data"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "node_count = 30\n",
      "edge_count = 20\n",
      "nodes = list(range(node_count))\n",
      "masses = np.random.random(node_count) * 20 + 10\n",
      "\n",
      "def build_edges():\n",
      "    edges = []\n",
      "    while len(edges) < edge_count:\n",
      "        sel_a = random.choice(nodes)\n",
      "        sel_b = random.choice(list(set(nodes) - set([sel_a])))\n",
      "        edges.append((sel_a, sel_b))\n",
      "        \n",
      "    return edges\n",
      "    \n",
      "edges = build_edges()\n",
      "\n",
      "# Initial nodes' position with normally distributed random numbers\n",
      "xpos = np.random.normal(size=node_count, scale=1000)\n",
      "ypos = np.random.normal(size=node_count, scale=1000)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Draw Graph"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "fig = bp.figure(width=600, height=600, x_range=(-500, 500), y_range=(-500, 500))\n",
      "\n",
      "fig.circle(xpos, ypos, radius=masses / 2, fill_alpha=0.6, name=\"nodes\")\n",
      "\n",
      "xlines = []\n",
      "ylines = []\n",
      "\n",
      "arr_edges = np.array(edges)\n",
      "xlines = xpos[arr_edges]\n",
      "ylines = ypos[arr_edges]\n",
      "fig.multi_line(xlines.tolist(), ylines.tolist(), name=\"edges\", line_alpha=0.5)\n",
      "\n",
      "fig.grid.grid_line_color = None\n",
      "fig.axis.axis_line_color = None\n",
      "fig.axis.major_tick_line_color = None"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# During developement, this remembers the original position of the graph.\n",
      "# So we can update the later cells only.\n",
      "orig_xpos = xpos.copy()\n",
      "orig_ypos = ypos.copy()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## The Layout Algorithm"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Reset to original (for interactively changing the notebook)\n",
      "xpos[:] = orig_xpos\n",
      "ypos[:] = orig_ypos\n",
      "\n",
      "DAMPENING = 0.2\n",
      "\n",
      "past_xpos = xpos.copy()\n",
      "past_ypos = ypos.copy()\n",
      "\n",
      "@jit\n",
      "def calc_force(i, j, fxs, fys, xs, ys, masses, strength):\n",
      "    dx = xs[i] - xs[j]\n",
      "    dy = ys[i] - ys[j]\n",
      "\n",
      "    dist = math.hypot(dx, dy)\n",
      "    theta = math.atan2(dy, dx)\n",
      "\n",
      "    optimal_dist = (masses[i] + masses[j])/2\n",
      "    force = (optimal_dist - dist) * strength\n",
      "\n",
      "    fx = math.cos(theta) * force\n",
      "    fy = math.sin(theta) * force\n",
      "    fxs[i] += fx\n",
      "    fys[i] += fy\n",
      "    fxs[j] -= fx\n",
      "    fys[j] -= fy\n",
      "\n",
      "    \n",
      "def update_force(fxs, fys, xs, ys, edges, masses):\n",
      "    for i, j in edges:\n",
      "        calc_force(i, j, fxs, fys, xs, ys, masses, 1.0)\n",
      "        \n",
      "    for i in range(len(xs)):\n",
      "        for j in range(i + 1, len(ys)):\n",
      "            calc_force(i, j, fxs, fys, xs, ys, masses, 1/len(xs))\n",
      "\n",
      "@jit\n",
      "def collision_avoid(xs, ys, masses):\n",
      "    for i in range(len(xs)):\n",
      "        for j in range(i + 1, len(ys)):\n",
      "            dx = xs[i] - xs[j]\n",
      "            dy = ys[i] - ys[j]\n",
      "            dist = math.hypot(dx, dy)\n",
      "            opt_dist = (masses[i] + masses[j])/1.8\n",
      "            if dist < opt_dist:\n",
      "                offset = (opt_dist - dist)/2\n",
      "                if abs(dx) < opt_dist:\n",
      "                    sign = -1 if dx < 0 else 1\n",
      "                    xs[i] += offset * sign\n",
      "                    xs[j] -= offset * sign\n",
      "                if abs(dy) < opt_dist:\n",
      "                    sign = -1 if dy < 0 else 1\n",
      "                    ys[i] += offset * sign\n",
      "                    ys[j] -= offset * sign\n",
      "\n",
      "                    \n",
      "def spring_fit_once(xs, ys, edges, masses, dt):\n",
      "    num = len(xs)\n",
      "    fxs = np.zeros(num, dtype=np.float32)\n",
      "    fys = np.zeros(num, dtype=np.float32)\n",
      "    update_force(fxs, fys, xs, ys, edges, masses)\n",
      "    dtdt = dt * dt\n",
      "    \n",
      "    # Mass = 1\n",
      "    axs = fxs #/masses\n",
      "    ays = fys #/masses\n",
      "    \n",
      "    # Verlet integration\n",
      "    new_xs = (2-DAMPENING) * xs - (1 - DAMPENING) * past_xpos + axs * dtdt\n",
      "    new_ys = (2-DAMPENING) * ys - (1 - DAMPENING) * past_ypos + ays * dtdt\n",
      "    \n",
      "    collision_avoid(new_xs, new_ys, masses)\n",
      "    \n",
      "    past_xpos[:] = xs\n",
      "    past_ypos[:] = ys\n",
      "    \n",
      "    xs[:] = new_xs\n",
      "    ys[:] = new_ys\n",
      "\n",
      "    \n",
      "renderer = fig.select(dict(name=\"nodes\", type=GlyphRenderer))\n",
      "ds_cir = renderer[0].data_source\n",
      "\n",
      "renderer = fig.select(dict(name=\"edges\", type=GlyphRenderer))\n",
      "ds_lines = renderer[0].data_source\n",
      "\n",
      "def update():\n",
      "    global xpos, ypos\n",
      "    max_width = np.max(masses)\n",
      "    spring_fit_once(xpos, ypos, edges, masses, dt=1/50)\n",
      "    ds_cir.data['x'] = xpos\n",
      "    ds_cir.data['y'] = ypos\n",
      "    ds_cir.push_notebook()\n",
      "    \n",
      "    xlines = xpos[arr_edges].tolist()\n",
      "    ylines = ypos[arr_edges].tolist()\n",
      "    ds_lines.data['xs'] = xlines\n",
      "    ds_lines.data['ys'] = ylines\n",
      "    ds_lines.push_notebook()\n",
      "    "
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Insert Plot"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "bp.show(fig)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Animate\n",
      "\n",
      "Run update for a fixed iteration count"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "for i in range(2000):\n",
      "    update()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}