{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "
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
" ] }, { "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 }