Numba implementation of tesseroid forward modeling: http://nbviewer.ipython.org/gist/leouieda/f2824b0062a4446a6e24

numba-tesseroids.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:65ca8db46ef1a2ab6ba3941f31f18d7d11c84a0091d9ec4cae2152861b1f2c88"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%matplotlib inline\n",
      "from __future__ import division, print_function\n",
      "import math\n",
      "import numpy as np\n",
      "import numpy.testing\n",
      "import multiprocessing\n",
      "import threading\n",
      "import numba\n",
      "from time import time\n",
      "from fatiando.mesher import TesseroidMesh, Tesseroid\n",
      "from fatiando import constants, utils, gridder\n",
      "from fatiando.gravmag.tesseroid import gz as gz_cython\n",
      "from fatiando.vis import mpl"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def make_data():    \n",
      "    model = TesseroidMesh((0, 5, 0, 10, 1000, -5000), (1, 5, 5))\n",
      "    model.addprop('density', 1000*np.ones(model.size))\n",
      "    x, y, z = gridder.regular((0, 5, 0, 10), (50, 50), z=50000)\n",
      "    return x, y, z, model\n",
      "\n",
      "def test_gz(func, **kwargs):\n",
      "    print(\"Testing {}:\".format(func.__name__))    \n",
      "    x, y, z, model = make_data()\n",
      "    # Test the func against the Cython version in Fatiando\n",
      "    f1 = func(x, y, z, model, **kwargs)\n",
      "    f2 = gz_cython(x, y, z, model)\n",
      "    np.testing.assert_allclose(f1, f2, err_msg=\"Results don't match current implementation\")\n",
      "    print(\" - Results match current implementation\")\n",
      "    # Time the execution\n",
      "    times = []\n",
      "    for i in xrange(5):\n",
      "        t1 = time()\n",
      "        func(x, y, z, model, **kwargs)\n",
      "        t2 = time()\n",
      "        times.append(t2 - t1)\n",
      "    print(\" - Execution time: {:f} s\".format(np.median(times)))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "nodes = np.array([-0.577350269189625731058868041146,\n",
      "                  0.577350269189625731058868041146])\n",
      "stack_size = 100"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def _check_input(lon, lat, height, ratio, out):\n",
      "    lon = np.radians(lon)\n",
      "    radius = height + constants.MEAN_EARTH_RADIUS\n",
      "    coslat = np.cos(np.radians(lat))\n",
      "    sinlat = np.sin(np.radians(lat))\n",
      "    if out is None:\n",
      "        result = np.zeros_like(lon)\n",
      "    else:\n",
      "        result = out        \n",
      "    return lon, coslat, sinlat, radius, ratio, result\n",
      "\n",
      "def _get_engine(engine):\n",
      "    if engine == 'numpy':\n",
      "        func = _numpy_engine\n",
      "    elif engine == 'numba':\n",
      "        func = _numba_engine\n",
      "    else:\n",
      "        raise ValueError(\"engine trouble\")\n",
      "    return func\n",
      "\n",
      "def _get_strides(size, njobs):    \n",
      "    n = size//njobs\n",
      "    strides = [(i*n, (i + 1)*n) for i in xrange(njobs - 1)]\n",
      "    strides.append((strides[-1][-1], size))\n",
      "    return strides\n",
      "\n",
      "def _split_arrays(arrays, args, njobs):\n",
      "    strides = _get_strides(arrays[0].size, njobs)\n",
      "    return [[x[low:high] for x in arrays] + args\n",
      "            for low, high in strides]\n",
      "\n",
      "def _dispatcher(args):\n",
      "    lon, coslat, sinlat, radius, result, model, engine, ratio = args  \n",
      "    func = _get_engine(engine)\n",
      "    for c in model:\n",
      "        density = c.props['density']\n",
      "        func(lon, coslat, sinlat, radius, c, density, ratio, result)\n",
      "    return result\n",
      "\n",
      "def gz(lon, lat, height, model, ratio=1.6, engine=None, njobs=None, out=None):\n",
      "    lon, coslat, sinlat, radius, ratio, result = _check_input(lon, lat, height, ratio, out)\n",
      "    if njobs is None or njobs == 1:\n",
      "        _dispatcher([lon, coslat, sinlat, radius, result, model, engine, ratio])\n",
      "    else:\n",
      "        chunks = _split_arrays([lon, coslat, sinlat, radius, result], [model, engine, ratio], njobs)\n",
      "        pool = multiprocessing.Pool(njobs)\n",
      "        result = np.hstack(pool.map(_dispatcher, chunks))\n",
      "        pool.close()\n",
      "    result *= constants.G*constants.SI2MGAL\n",
      "    return result"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def _numpy_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):\n",
      "    for l in xrange(lon.size):\n",
      "        stack = [cell]\n",
      "        while stack:\n",
      "            t = stack.pop()\n",
      "            distance, Llon, Llat, Lr = distance_size_numpy(lon[l], coslat[l], sinlat[l], radius[l], t)\n",
      "            nlon = 1 if distance/Llon > ratio else 2\n",
      "            nlat = 1 if distance/Llat > ratio else 2\n",
      "            nr = 1 if distance/Lr > ratio else 2\n",
      "            new_cells = nlon*nlat*nr\n",
      "            if new_cells > 1:\n",
      "                if new_cells + len(stack) > stack_size:\n",
      "                    print(distance, Llon, Llat, Lr)\n",
      "                    raise OverflowError('Tesseroid stack overflowed')\n",
      "                stack.extend(t.split(nlon, nlat, nr))\n",
      "            else:\n",
      "                lonc, sinlatc, coslatc, rc, scale = scale_nodes_numpy(t.get_bounds(), nodes)\n",
      "                result[l] += scale*density*kernelz_numpy(lon[l], coslat[l], sinlat[l], radius[l], \n",
      "                                                         lonc, sinlatc, coslatc, rc)\n",
      "\n",
      "def scale_nodes_numpy(bounds, nodes):\n",
      "    w, e, s, n, top, bottom = bounds\n",
      "    d2r = np.pi/180\n",
      "    dlon = d2r*(e - w)\n",
      "    dlat = d2r*(n - s)\n",
      "    dr = top - bottom\n",
      "    # Scale the GLQ nodes to the integration limits\n",
      "    lonc = 0.5*dlon*nodes + d2r*0.5*(e + w)\n",
      "    latc = 0.5*dlat*nodes + d2r*0.5*(n + s)\n",
      "    sinlatc = np.sin(latc)\n",
      "    coslatc = np.cos(latc)\n",
      "    rc = (0.5*dr*nodes +\n",
      "          0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))\n",
      "    scale = dlon*dlat*dr*0.125\n",
      "    return lonc, sinlatc, coslatc, rc, scale\n",
      "    \n",
      "def distance_size_numpy(lon, coslat, sinlat, radius, cell):\n",
      "    w, e, s, n, top, bottom = cell.get_bounds()\n",
      "    d2r = np.pi/180\n",
      "    rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS\n",
      "    lont = d2r*0.5*(w + e)\n",
      "    latt = d2r*0.5*(s + n)\n",
      "    sinlatt = np.sin(latt)\n",
      "    coslatt = np.cos(latt)\n",
      "    cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)\n",
      "    distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)\n",
      "    # Calculate the dimensions of the tesseroid in meters\n",
      "    rtop = top + constants.MEAN_EARTH_RADIUS\n",
      "    Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))\n",
      "    Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))\n",
      "    Lr = top - bottom\n",
      "    return distance, Llon, Llat, Lr\n",
      "\n",
      "def kernelz_numpy(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):\n",
      "    r_sqr = radius**2\n",
      "    result = 0\n",
      "    for i in range(2):\n",
      "        coslon = np.cos(lon - lonc[i])\n",
      "        for j in range(2):\n",
      "            cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon\n",
      "            for k in range(2):\n",
      "                l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi\n",
      "                kappa = (rc[k]**2)*coslatc[j]\n",
      "                result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)\n",
      "    result *= -1\n",
      "    return result"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "@numba.jit(looplift=True, nogil=True)\n",
      "def _numba_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):\n",
      "    bounds = np.array(cell.get_bounds())\n",
      "    stack = np.empty((stack_size, 6))\n",
      "    lonc = np.empty_like(nodes)\n",
      "    sinlatc = np.empty_like(nodes)\n",
      "    coslatc = np.empty_like(nodes)\n",
      "    rc = np.empty_like(nodes)\n",
      "    for l in range(result.size):\n",
      "        for i in range(6):\n",
      "            stack[0, i] = bounds[i]\n",
      "        stktop = 0\n",
      "        while stktop >= 0:\n",
      "            w, e, s, n, top, bottom = stack[stktop, :]\n",
      "            stktop -= 1\n",
      "            distance, Llon, Llat, Lr = distance_size_numba(lon[l], coslat[l], sinlat[l], radius[l], \n",
      "                                                           w, e, s, n, top, bottom)\n",
      "            nlon, nlat, nr, new_cells = divisions_per_dimension(distance, Llon, Llat, Lr, ratio)\n",
      "            if new_cells > 1:\n",
      "                if new_cells + (stktop + 1) > stack_size:\n",
      "                    raise OverflowError\n",
      "                stktop = split_numba(w, e, s, n, top, bottom, nlon, nlat, nr,\n",
      "                                     stack, stktop)\n",
      "            else:\n",
      "                scale = scale_nodes_numba(w, e, s, n, top, bottom, nodes, \n",
      "                                          lonc, sinlatc, coslatc, rc)\n",
      "                result[l] += density*scale*kernelz_numba(lon[l], coslat[l], sinlat[l], radius[l], \n",
      "                                                lonc, sinlatc, coslatc, rc)\n",
      "    \n",
      "@numba.jit(nopython=True, nogil=True)\n",
      "def scale_nodes_numba(w, e, s, n, top, bottom, nodes, lonc, sinlatc, coslatc, rc):\n",
      "    d2r = np.pi/180\n",
      "    dlon = d2r*(e - w)\n",
      "    dlat = d2r*(n - s)\n",
      "    dr = top - bottom\n",
      "    # Scale the GLQ nodes to the integration limits\n",
      "    for i in range(len(nodes)):\n",
      "        lonc[i] = 0.5*dlon*nodes[i] + d2r*0.5*(e + w)\n",
      "        latc = 0.5*dlat*nodes[i] + d2r*0.5*(n + s)\n",
      "        sinlatc[i] = np.sin(latc)\n",
      "        coslatc[i] = np.cos(latc)\n",
      "        rc[i] = (0.5*dr*nodes[i] +\n",
      "              0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))\n",
      "    scale = dlon*dlat*dr*0.125\n",
      "    return scale\n",
      "    \n",
      "@numba.jit(nopython=True, nogil=True)\n",
      "def distance_size_numba(lon, coslat, sinlat, radius, w, e, s, n, top, bottom):\n",
      "    d2r = np.pi/180\n",
      "    rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS\n",
      "    lont = d2r*0.5*(w + e)\n",
      "    latt = d2r*0.5*(s + n)\n",
      "    sinlatt = np.sin(latt)\n",
      "    coslatt = np.cos(latt)\n",
      "    cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)\n",
      "    distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)\n",
      "    # Calculate the dimensions of the tesseroid in meters\n",
      "    rtop = top + constants.MEAN_EARTH_RADIUS\n",
      "    Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))\n",
      "    Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))\n",
      "    Lr = top - bottom\n",
      "    return distance, Llon, Llat, Lr\n",
      "\n",
      "@numba.jit(nopython=True, nogil=True)\n",
      "def kernelz_numba(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):\n",
      "    r_sqr = radius**2\n",
      "    result = 0\n",
      "    for i in range(2):\n",
      "        coslon = np.cos(lon - lonc[i])\n",
      "        for j in range(2):\n",
      "            cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon\n",
      "            for k in range(2):\n",
      "                l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi\n",
      "                kappa = (rc[k]**2)*coslatc[j]\n",
      "                result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)\n",
      "    result *= -1\n",
      "    return result\n",
      "\n",
      "@numba.jit(nopython=True, nogil=True)\n",
      "def split_numba(w, e, s, n, top, bottom, nlon, nlat, nr, stack, stktop):\n",
      "    dlon = (e - w)/nlon\n",
      "    dlat = (n - s)/nlat\n",
      "    dr = (top - bottom)/nr\n",
      "    for i in xrange(nlon):\n",
      "        for j in xrange(nlat):\n",
      "            for k in xrange(nr):\n",
      "                stktop += 1\n",
      "                stack[stktop, 0] = w + i*dlon\n",
      "                stack[stktop, 1] = w + (i + 1)*dlon\n",
      "                stack[stktop, 2] = s + j*dlat\n",
      "                stack[stktop, 3] = s + (j + 1)*dlat\n",
      "                stack[stktop, 4] = bottom + (k + 1)*dr\n",
      "                stack[stktop, 5] = bottom + k*dr\n",
      "    return stktop\n",
      "\n",
      "@numba.jit(nopython=True, nogil=True)\n",
      "def divisions_per_dimension(distance, Llon, Llat, Lr, ratio):\n",
      "    nlon = 1 if distance/Llon > ratio else 2\n",
      "    nlat = 1 if distance/Llat > ratio else 2\n",
      "    nr = 1 if distance/Lr > ratio else 2\n",
      "    return nlon, nlat, nr, nlon*nlat*nr"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "test_gz(gz_cython)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Testing gz:\n",
        " - Results match current implementation"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        " - Execution time: 0.259152 s"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n"
       ]
      }
     ],
     "prompt_number": 7
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "test_gz(gz, engine='numba')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Testing gz:\n",
        " - Results match current implementation"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        " - Execution time: 0.212037 s"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n"
       ]
      }
     ],
     "prompt_number": 8
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "test_gz(gz, engine='numpy')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Testing gz:\n",
        " - Results match current implementation"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n",
        " - Execution time: 13.163021 s"
       ]
      },
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "\n"
       ]
      }
     ],
     "prompt_number": 9
    }
   ],
   "metadata": {}
  }
 ]
}