poliastro benchmark

poliastro_benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datetime import timedelta\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from astropy import units as u\n",
    "from astropy.coordinates import CartesianDifferential, CartesianRepresentation\n",
    "from astropy.time import Time\n",
    "\n",
    "from poliastro.bodies import Sun\n",
    "from poliastro.frames import Planes\n",
    "from poliastro.twobody import Orbit\n",
    "from poliastro.twobody.propagation import propagate\n",
    "\n",
    "from poliastro.core._jit import jit\n",
    "from poliastro.core.elements import coe2rv\n",
    "from poliastro.core.propagation.farnocchia import delta_t_from_nu, nu_from_delta_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DAYS = 36525\n",
    "\n",
    "ASTRO_TIME_ZERO = Time('1970-01-01 00:00')\n",
    "\n",
    "MPC_RAW = \"\"\"{\n",
    "\"NEO_flag\": 1,\n",
    "\"One_km_NEO_flag\": 1,\n",
    "\"H\": 11.16,\n",
    "\"G\": 0.46,\n",
    "\"Num_obs\": 8406,\n",
    "\"rms\": 0.69,\n",
    "\"U\": \"0\",\n",
    "\"Arc_years\": \"1893-2019\",\n",
    "\"Perturbers\": \"M-v\",\n",
    "\"Perturbers_2\": \"38h\",\n",
    "\"Number\": \"(433)\",\n",
    "\"Name\": \"Eros\",\n",
    "\"Principal_desig\": \"A898 PA\",\n",
    "\"Other_desigs\": [\n",
    "\"1956 PC\"\n",
    "],\n",
    "\"Epoch\": 2458700.5,\n",
    "\"M\": 103.1926,\n",
    "\"Peri\": 178.84531,\n",
    "\"Node\": 304.30277,\n",
    "\"i\": 10.82872,\n",
    "\"e\": 0.2227412,\n",
    "\"n\": 0.55971324,\n",
    "\"a\": 1.4582288,\n",
    "\"Ref\": \"MPO467611\",\n",
    "\"Num_opps\": 52,\n",
    "\"Computer\": \"MPCLINUX\",\n",
    "\"Hex_flags\": \"1804\",\n",
    "\"Last_obs\": \"2019-05-09\",\n",
    "\"Tp\": 2458516.13308,\n",
    "\"Orbital_period\": 1.7609155,\n",
    "\"Perihelion_dist\": 1.1334212,\n",
    "\"Aphelion_dist\": 1.7830364,\n",
    "\"Semilatus_rectum\": 0.6929404,\n",
    "\"Synodic_period\": 2.3142064,\n",
    "\"Orbit_type\": \"Amor\"\n",
    "}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def orbit_from_mpc(data):\n",
    "    return Orbit.from_classical(\n",
    "        Sun,\n",
    "        a = (data['a'] * u.AU).to(u.km),\n",
    "        ecc = data['e'] * u.one,\n",
    "        inc = (data['i'] * u.deg).to(u.rad),\n",
    "        raan = (data['Node'] * u.deg).to(u.rad),\n",
    "        argp = (data['Peri'] * u.deg).to(u.rad),\n",
    "        nu = (data['M'] * u.deg).to(u.rad),\n",
    "        epoch = Time(data[\"Epoch\"], format = 'jd'),\n",
    "        plane = Planes.EARTH_ECLIPTIC,\n",
    "    )\n",
    "\n",
    "mpc_data = json.loads(MPC_RAW)\n",
    "mpc_orb = orbit_from_mpc(mpc_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: ErfaWarning: ERFA function \"dtf2d\" yielded 1 of \"dubious year (Note 6)\" [astropy._erfa.core]\n"
     ]
    }
   ],
   "source": [
    "time_zero = ASTRO_TIME_ZERO.datetime\n",
    "astro_times = Time([Time(time_zero + timedelta(days = day)) for day in range(0, DAYS, 1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 527 ms, sys: 3.92 ms, total: 531 ms\n",
      "Wall time: 529 ms\n"
     ]
    }
   ],
   "source": [
    "eph_A = np.zeros((astro_times.shape[0], 3), dtype = 'f4')\n",
    "\n",
    "def test_A():\n",
    "    relative_times = astro_times - mpc_orb.epoch\n",
    "    eph_A[:, :] = np.transpose(propagate(mpc_orb, relative_times).xyz.to(u.AU).value)\n",
    "\n",
    "%time test_A()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def farnocchia_coe_fast(k, p, ecc, inc, raan, argp, nu0, tof):\n",
    "    \n",
    "    q = p / (1 + ecc)\n",
    "\n",
    "    delta_t0 = delta_t_from_nu(nu0, ecc, k, q)\n",
    "    delta_t = delta_t0 + tof\n",
    "\n",
    "    nu = nu_from_delta_t(delta_t, ecc, k, q)\n",
    "\n",
    "    return coe2rv(k, p, ecc, inc, raan, argp, nu)\n",
    "\n",
    "def farnocchia_coe(k, p, ecc, inc, raan, argp, nu0, tofs):\n",
    "    \n",
    "    k = k.to(u.km ** 3 / u.s ** 2).value\n",
    "    p = p.to_value(u.km)\n",
    "    ecc = ecc.value\n",
    "    inc = inc.to_value(u.rad)\n",
    "    raan = raan.to_value(u.rad)\n",
    "    argp = argp.to_value(u.rad)\n",
    "    nu0 = nu0.to_value(u.rad)\n",
    "    tofs = tofs.to(u.s).value\n",
    "    \n",
    "    results = [farnocchia_coe_fast(k, p, ecc, inc, raan, argp, nu0, tof) for tof in tofs]\n",
    "    \n",
    "    # TODO: Rewrite to avoid iterating twice\n",
    "    return (\n",
    "        [result[0] for result in results] * u.km,\n",
    "        [result[1] for result in results] * u.km / u.s,\n",
    "    )\n",
    "\n",
    "def propagate_fast(orbit, time_of_flight, *, method=farnocchia_coe, rtol=1e-10, **kwargs):\n",
    "\n",
    "    rr, vv = method(\n",
    "        orbit.attractor.k,\n",
    "        orbit.p,\n",
    "        orbit.ecc,\n",
    "        orbit.inc,\n",
    "        orbit.raan,\n",
    "        orbit.argp,\n",
    "        orbit.nu,\n",
    "        time_of_flight.reshape(-1).to(u.s),\n",
    "    )\n",
    "    \n",
    "    # TODO: Turn these into unit tests\n",
    "    assert rr.ndim == 2\n",
    "    assert vv.ndim == 2\n",
    "\n",
    "    cartesian = CartesianRepresentation(\n",
    "        rr, differentials = CartesianDifferential(vv, xyz_axis = 1), xyz_axis = 1\n",
    "    )\n",
    "\n",
    "    return cartesian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 388 ms, sys: 3.99 ms, total: 392 ms\n",
      "Wall time: 393 ms\n"
     ]
    }
   ],
   "source": [
    "eph_B = np.zeros((astro_times.shape[0], 3), dtype = 'f4')\n",
    "\n",
    "def test_B():\n",
    "    relative_times = astro_times - mpc_orb.epoch\n",
    "    eph_B[:, :] = np.transpose(propagate_fast(mpc_orb, relative_times).xyz.to(u.AU).value)\n",
    "\n",
    "%time test_B()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.all((eph_B - eph_A) == 0.0) # not even floating point inaccuracies ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}