{
 "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",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show, output_notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 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 list(range((max(NBX,NBY))//2)):\n",
    "        for xi in list(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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xs, ys = streamlines(xx, yy, U.T, V.T, density=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "p = figure()\n",
    "p.segment(x0, y0, x1, y1, line_color=colors, line_width=2,name=\"vector example\")\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "q = figure()\n",
    "q.multi_line(xs, ys, line_color=\"#ee6666\", line_width=2, line_alpha=0.8)\n",
    "show(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
