{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "raw",
     "metadata": {},
     "source": [
      "This IPython Notebook contains simple examples of the line function. \n",
      "\n",
      "To clear all previously rendered cell outputs, select from the menu:\n",
      "\n",
      "    Cell -> All Output -> Clear"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "from scipy.integrate import odeint"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from bokeh.plotting import *"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# adapted from here: http://www.atm.damtp.cam.ac.uk/people/tjf37/streamplot.py\n",
      "# using until Bokeh has its own streamline integrator built in\n",
      "def streamlines(x, y, u, v, density=1):\n",
      "    '''Returns streamlines of a vector flow.\n",
      "\n",
      "    * x and y are 1d arrays defining an *evenly spaced* grid.\n",
      "    * u and v are 2d arrays (shape [y,x]) giving velocities.\n",
      "    * density controls the closeness of the streamlines. For different\n",
      "      densities in each direction, use a tuple or list [densityx, densityy].\n",
      "\n",
      "     '''\n",
      "\n",
      "    ## Set up some constants - size of the grid used.\n",
      "    NGX = len(x)\n",
      "    NGY = len(y)\n",
      "\n",
      "    ## Constants used to convert between grid index coords and user coords.\n",
      "    DX = x[1]-x[0]\n",
      "    DY = y[1]-y[0]\n",
      "    XOFF = x[0]\n",
      "    YOFF = y[0]\n",
      "\n",
      "    ## Now rescale velocity onto axes-coordinates\n",
      "    u = u / (x[-1]-x[0])\n",
      "    v = v / (y[-1]-y[0])\n",
      "    speed = np.sqrt(u*u+v*v)\n",
      "    ## s (path length) will now be in axes-coordinates, but we must\n",
      "    ## rescale u for integrations.\n",
      "    u *= NGX\n",
      "    v *= NGY\n",
      "    ## Now u and v in grid-coordinates.\n",
      "\n",
      "    NBX = int(30*density)\n",
      "    NBY = int(30*density)\n",
      "    blank = np.zeros((NBY,NBX))\n",
      "\n",
      "    bx_spacing = NGX/float(NBX-1)\n",
      "    by_spacing = NGY/float(NBY-1)\n",
      "\n",
      "    def blank_pos(xi, yi):\n",
      "        return int((xi / bx_spacing) + 0.5), \\\n",
      "               int((yi / by_spacing) + 0.5)\n",
      "\n",
      "    def value_at(a, xi, yi):\n",
      "        if type(xi) == np.ndarray:\n",
      "            x = xi.astype(np.int)\n",
      "            y = yi.astype(np.int)\n",
      "        else:\n",
      "            x = np.int(xi)\n",
      "            y = np.int(yi)\n",
      "        a00 = a[y,x]\n",
      "        a01 = a[y,x+1]\n",
      "        a10 = a[y+1,x]\n",
      "        a11 = a[y+1,x+1]\n",
      "        xt = xi - x\n",
      "        yt = yi - y\n",
      "        a0 = a00*(1-xt) + a01*xt\n",
      "        a1 = a10*(1-xt) + a11*xt\n",
      "        return a0*(1-yt) + a1*yt\n",
      "\n",
      "    def rk4_integrate(x0, y0):\n",
      "        ## This function does RK4 forward and back trajectories from\n",
      "        ## the initial conditions, with the odd 'blank array'\n",
      "        ## termination conditions. TODO tidy the integration loops.\n",
      "\n",
      "        def f(xi, yi):\n",
      "            dt_ds = 1./value_at(speed, xi, yi)\n",
      "            ui = value_at(u, xi, yi)\n",
      "            vi = value_at(v, xi, yi)\n",
      "            return ui*dt_ds, vi*dt_ds\n",
      "\n",
      "        def g(xi, yi):\n",
      "            dt_ds = 1./value_at(speed, xi, yi)\n",
      "            ui = value_at(u, xi, yi)\n",
      "            vi = value_at(v, xi, yi)\n",
      "            return -ui*dt_ds, -vi*dt_ds\n",
      "\n",
      "        check = lambda xi, yi: xi>=0 and xi<NGX-1 and yi>=0 and yi<NGY-1\n",
      "\n",
      "        bx_changes = []\n",
      "        by_changes = []\n",
      "\n",
      "        ## Integrator function\n",
      "        def rk4(x0, y0, f):\n",
      "            ds = 0.01 #min(1./NGX, 1./NGY, 0.01)\n",
      "            stotal = 0\n",
      "            xi = x0\n",
      "            yi = y0\n",
      "            xb, yb = blank_pos(xi, yi)\n",
      "            xf_traj = []\n",
      "            yf_traj = []\n",
      "            while check(xi, yi):\n",
      "                # Time step. First save the point.\n",
      "                xf_traj.append(xi)\n",
      "                yf_traj.append(yi)\n",
      "                # Next, advance one using RK4\n",
      "                try:\n",
      "                    k1x, k1y = f(xi, yi)\n",
      "                    k2x, k2y = f(xi + .5*ds*k1x, yi + .5*ds*k1y)\n",
      "                    k3x, k3y = f(xi + .5*ds*k2x, yi + .5*ds*k2y)\n",
      "                    k4x, k4y = f(xi + ds*k3x, yi + ds*k3y)\n",
      "                except IndexError:\n",
      "                    # Out of the domain on one of the intermediate steps\n",
      "                    break\n",
      "                xi += ds*(k1x+2*k2x+2*k3x+k4x) / 6.\n",
      "                yi += ds*(k1y+2*k2y+2*k3y+k4y) / 6.\n",
      "                # Final position might be out of the domain\n",
      "                if not check(xi, yi): break\n",
      "                stotal += ds\n",
      "                # Next, if s gets to thres, check blank.\n",
      "                new_xb, new_yb = blank_pos(xi, yi)\n",
      "                if new_xb != xb or new_yb != yb:\n",
      "                    # New square, so check and colour. Quit if required.\n",
      "                    if blank[new_yb,new_xb] == 0:\n",
      "                        blank[new_yb,new_xb] = 1\n",
      "                        bx_changes.append(new_xb)\n",
      "                        by_changes.append(new_yb)\n",
      "                        xb = new_xb\n",
      "                        yb = new_yb\n",
      "                    else:\n",
      "                        break\n",
      "                if stotal > 2:\n",
      "                    break\n",
      "            return stotal, xf_traj, yf_traj\n",
      "\n",
      "        integrator = rk4\n",
      "\n",
      "        sf, xf_traj, yf_traj = integrator(x0, y0, f)\n",
      "        sb, xb_traj, yb_traj = integrator(x0, y0, g)\n",
      "        stotal = sf + sb\n",
      "        x_traj = xb_traj[::-1] + xf_traj[1:]\n",
      "        y_traj = yb_traj[::-1] + yf_traj[1:]\n",
      "\n",
      "        ## Tests to check length of traj. Remember, s in units of axes.\n",
      "        if len(x_traj) < 1: return None\n",
      "        if stotal > .2:\n",
      "            initxb, inityb = blank_pos(x0, y0)\n",
      "            blank[inityb, initxb] = 1\n",
      "            return x_traj, y_traj\n",
      "        else:\n",
      "            for xb, yb in zip(bx_changes, by_changes):\n",
      "                blank[yb, xb] = 0\n",
      "            return None\n",
      "\n",
      "    ## A quick function for integrating trajectories if blank==0.\n",
      "    trajectories = []\n",
      "    def traj(xb, yb):\n",
      "        if xb < 0 or xb >= NBX or yb < 0 or yb >= NBY:\n",
      "            return\n",
      "        if blank[yb, xb] == 0:\n",
      "            t = rk4_integrate(xb*bx_spacing, yb*by_spacing)\n",
      "            if t != None:\n",
      "                trajectories.append(t)\n",
      "\n",
      "    ## Now we build up the trajectory set. I've found it best to look\n",
      "    ## for blank==0 along the edges first, and work inwards.\n",
      "    for indent in range((max(NBX,NBY))/2):\n",
      "        for xi in range(max(NBX,NBY)-2*indent):\n",
      "            traj(xi+indent, indent)\n",
      "            traj(xi+indent, NBY-1-indent)\n",
      "            traj(indent, xi+indent)\n",
      "            traj(NBX-1-indent, xi+indent)\n",
      "\n",
      "    xs = [np.array(t[0])*DX+XOFF for t in trajectories]\n",
      "    ys = [np.array(t[1])*DY+YOFF for t in trajectories]\n",
      "\n",
      "    return xs, ys"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "xx = np.linspace(-3, 3, 100)\n",
      "yy = np.linspace(-3, 3, 100)\n",
      "\n",
      "Y, X = np.meshgrid(xx, yy)\n",
      "U = -1 - X**2 + Y\n",
      "V = 1 + X - Y**2\n",
      "speed = np.sqrt(U*U + V*V)\n",
      "theta = np.arctan(V/U)\n",
      "\n",
      "x0 = X[::2, ::2].flatten()\n",
      "y0 = Y[::2, ::2].flatten()\n",
      "length = speed[::2, ::2].flatten()/40\n",
      "angle = theta[::2, ::2].flatten()\n",
      "x1 = x0 + length * np.cos(angle)\n",
      "y1 = y0 + length * np.sin(angle)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "xs, ys = streamlines(xx, yy, U.T, V.T, density=2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "cm = np.array([\"#C7E9B4\", \"#7FCDBB\", \"#41B6C4\", \"#1D91C0\", \"#225EA8\", \"#0C2C84\"])\n",
      "ix = ((length-length.min())/(length.max()-length.min())*5).astype('int')\n",
      "colors = cm[ix]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "output_notebook()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "figure(tools=\"pan,wheel_zoom,box_zoom,reset,previewsave\")"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "segment(x0, y0, x1, y1, line_color=colors, line_width=2,name=\"vector example\")\n",
      "show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "multi_line(xs, ys, line_color=\"#ee6666\", line_width=2, line_alpha=0.8)\n",
      "show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}