examples/Solving Lambert's problem in Python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">For executing this notebook you will need to use the branch `iod-numba` of the poliastro repository: https://github.com/Pybonacci/poliastro/tree/iod-numba</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from astropy import units as u\n",
    "\n",
    "from poliastro.bodies import Earth\n",
    "from poliastro.iod import lambert\n",
    "from poliastro.math import dot\n",
    "from poliastro.stumpff import c2, c3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "k = Earth.k\n",
    "r0 = [15945.34, 0.0, 0.0] * u.km\n",
    "r = [12214.83399, 10249.46731, 0.0] * u.km\n",
    "tof = 76.0 * u.min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.05891159,  2.91596286,  0.        ]),\n",
       " array([-3.45156464,  0.91031284,  0.        ]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lambert(k.to(u.km ** 3 / u.s ** 2).value,\n",
    "        r0.value, r.value,\n",
    "        tof.to(u.s).value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare inputs to measure performance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "k_ = k.to(u.km ** 3 / u.s ** 2).value\n",
    "r0_ = r0.to(u.km).value\n",
    "r_ = r.to(u.km).value\n",
    "tof_ = tof.to(u.s).value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The slowest run took 7.36 times longer than the fastest. This could mean that an intermediate result is being cached \n",
      "100000 loops, best of 3: 14.9 µs per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit lambert(k_, r0_, r_, tof_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's implement the simplest one: Bate-Mueller-White universal variable approach, with a bisection method as suggested by Vallado."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def lambert_py(k, r0, r, tof, short=True, numiter=35, rtol=1e-6):\n",
    "    if short:\n",
    "        t_m = +1\n",
    "    else:\n",
    "        t_m = -1\n",
    "\n",
    "    norm_r0 = np.dot(r0, r0)**.5\n",
    "    norm_r = np.dot(r, r)**.5\n",
    "    cos_dnu = np.dot(r0, r) / (norm_r0 * norm_r)\n",
    "    sin_dnu = t_m * (1 - cos_dnu ** 2)**.5\n",
    "\n",
    "    A = t_m * (norm_r * norm_r0 * (1 + cos_dnu))**.5\n",
    "\n",
    "    if A == 0.0:\n",
    "        raise RuntimeError(\"Cannot compute orbit\")\n",
    "\n",
    "    psi = 0.0\n",
    "    psi_low = -4 * np.pi\n",
    "    psi_up = 4 * np.pi\n",
    "\n",
    "    count = 0\n",
    "    while count < numiter:\n",
    "        y = norm_r0 + norm_r + A * (psi * c3(psi) - 1) / c2(psi)**.5\n",
    "        if A > 0.0 and y < 0.0:\n",
    "            # Readjust xi_low until y > 0.0 (?)\n",
    "            pass\n",
    "        xi = np.sqrt(y / c2(psi))\n",
    "        tof_new = (xi**3 * c3(psi) + A * np.sqrt(y)) / np.sqrt(k)\n",
    "\n",
    "        # Convergence check\n",
    "        if np.abs((tof_new - tof) / tof) < rtol:\n",
    "            break\n",
    "        else:\n",
    "            count += 1\n",
    "            # Bisection check\n",
    "            if tof_new <= tof:\n",
    "                psi_low = psi\n",
    "            else:\n",
    "                psi_up = psi\n",
    "            psi = (psi_up + psi_low) / 2\n",
    "    else:\n",
    "        raise RuntimeError(\"Convergence could not be achieved under \"\n",
    "                           \"%d iterations\" % numiter)\n",
    "\n",
    "    f = 1 - y / norm_r0\n",
    "    g = A * np.sqrt(y / k)\n",
    "\n",
    "    gdot = 1 - y / norm_r\n",
    "\n",
    "    v0 = (r - f * r0) / g\n",
    "    v = (gdot * r - r0) / g\n",
    "\n",
    "    return v0, v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.05890855,  2.91596498,  0.        ]),\n",
       " array([-3.45156368,  0.91031642,  0.        ]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lambert_py(k_, r0_, r_, tof_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000 loops, best of 3: 222 µs per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit lambert_py(k_, r0_, r_, tof_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numba\n",
    "\n",
    "def lambert_numba(k, r0, r, tof, short=True, numiter=35, rtol=1e-8):\n",
    "    try:\n",
    "        f, g, fdot, gdot = _lambert(k, r0, r, tof, short, numiter, rtol)\n",
    "    except RuntimeError as e:\n",
    "        raise e\n",
    "\n",
    "    v0 = (r - f * r0) / g\n",
    "    v = (gdot * r - r0) / g\n",
    "\n",
    "    return v0, v\n",
    "\n",
    "@numba.njit\n",
    "def _lambert(k, r0, r, tof, short, numiter, rtol):\n",
    "    if short:\n",
    "        t_m = +1\n",
    "    else:\n",
    "        t_m = -1\n",
    "\n",
    "    norm_r0 = dot(r0, r0)**.5\n",
    "    norm_r = dot(r, r)**.5\n",
    "    cos_dnu = dot(r0, r) / (norm_r0 * norm_r)\n",
    "    sin_dnu = t_m * (1 - cos_dnu ** 2)**.5\n",
    "\n",
    "    A = t_m * (norm_r * norm_r0 * (1 + cos_dnu))**.5\n",
    "\n",
    "    if A == 0.0:\n",
    "        raise RuntimeError\n",
    "\n",
    "    psi = 0.0\n",
    "    psi_low = -4 * np.pi\n",
    "    psi_up = 4 * np.pi\n",
    "\n",
    "    count = 0\n",
    "    while count < numiter:\n",
    "        y = norm_r0 + norm_r + A * (psi * c3(psi) - 1) / c2(psi)**.5\n",
    "        if A > 0.0 and y < 0.0:\n",
    "            # Readjust xi_low until y > 0.0\n",
    "            # Translated directly from Vallado\n",
    "            while y < 0.0:\n",
    "                psi_low = psi\n",
    "                psi = 0.8 * (1.0 / c3(psi)) * (1.0 - (norm_r0 + norm_r) * np.sqrt(c2(psi)) / A)\n",
    "                y = norm_r0 + norm_r + A * (psi * c3(psi) - 1) / c2(psi)**.5\n",
    "\n",
    "        xi = np.sqrt(y / c2(psi))\n",
    "        tof_new = (xi**3 * c3(psi) + A * np.sqrt(y)) / np.sqrt(k)\n",
    "\n",
    "        # Convergence check\n",
    "        if np.abs((tof_new - tof) / tof) < rtol:\n",
    "            break\n",
    "        else:\n",
    "            count += 1\n",
    "            # Bisection check\n",
    "            if tof_new <= tof:\n",
    "                psi_low = psi\n",
    "            else:\n",
    "                psi_up = psi\n",
    "            psi = (psi_up + psi_low) / 2\n",
    "    else:\n",
    "        raise RuntimeError\n",
    "\n",
    "    f = 1 - y / norm_r0\n",
    "    g = A * np.sqrt(y / k)\n",
    "\n",
    "    gdot = 1 - y / norm_r\n",
    "\n",
    "    return f, g, (f * gdot - 1) / g, gdot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.05890993,  2.91596401,  0.        ]),\n",
       " array([-3.45156412,  0.91031479,  0.        ]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lambert_numba(k_, r0_, r_, tof_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The slowest run took 6.45 times longer than the fastest. This could mean that an intermediate result is being cached \n",
      "10000 loops, best of 3: 21.6 µs per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit lambert_numba(k_, r0_, r_, tof_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================= test session starts ==============================\n",
      "platform linux -- Python 3.4.3 -- py-1.4.26 -- pytest-2.6.4\n",
      "collected 39 items\n",
      "\n",
      "../poliastro/tests/test_bodies.py ....\n",
      "../poliastro/tests/test_iod.py ...\n",
      "../poliastro/tests/test_maneuver.py ......\n",
      "../poliastro/tests/test_math.py ....\n",
      "../poliastro/tests/test_plotting.py ..\n",
      "../poliastro/tests/test_twobody.py .................\n",
      "../poliastro/tests/test_util.py ...\n",
      "\n",
      "========================== 39 passed in 0.74 seconds ===========================\n"
     ]
    }
   ],
   "source": [
    "import poliastro\n",
    "import poliastro.iod\n",
    "poliastro.iod.lambert = lambert_numba\n",
    "\n",
    "poliastro.test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}