propagate_02.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "FN = 'nea_extended.json'\n",
    "\n",
    "from bulk_lambert.bulk_many import propagate_many\n",
    "\n",
    "import json\n",
    "\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "from astropy.time import Time\n",
    "from astropy import units as u\n",
    "\n",
    "from poliastro.bodies import Sun\n",
    "from poliastro.frames import Planes\n",
    "from poliastro.twobody import Orbit\n",
    "\n",
    "from orbital.utilities import true_anomaly_from_mean\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def orbit_from_mpc(body):\n",
    "\n",
    "    nu = true_anomaly_from_mean(\n",
    "        e = body['e'],\n",
    "        M = float((body['M'] * u.deg).to(u.rad).value)\n",
    "    ) * u.rad\n",
    "\n",
    "    if not -np.pi * u.rad <= nu < np.pi * u.rad:\n",
    "        nu = ((nu + np.pi * u.rad) % (2 * np.pi * u.rad) - np.pi * u.rad).to(\n",
    "            nu.unit\n",
    "        )\n",
    "\n",
    "    return Orbit.from_classical(\n",
    "        Sun,\n",
    "        a = (body['a'] * u.AU).to(u.m),\n",
    "        ecc = body['e'] * u.one,\n",
    "        inc = (body['i'] * u.deg).to(u.rad),\n",
    "        raan = (body['Node'] * u.deg).to(u.rad),\n",
    "        argp = (body['Peri'] * u.deg).to(u.rad),\n",
    "        nu = nu,\n",
    "        epoch = Time(body[\"Epoch\"], format = 'jd'),\n",
    "        plane = Planes.EARTH_ECLIPTIC,\n",
    "    )\n",
    "\n",
    "def read_mpc(fn):\n",
    "\n",
    "    with open(fn, 'r') as f:\n",
    "        raw = json.load(f)\n",
    "        \n",
    "    # return [orbit_from_mpc(body) for body in raw]\n",
    "    return Parallel(n_jobs=8)(delayed(orbit_from_mpc)(body) for body in raw)\n",
    "\n",
    "mpc_orbits = read_mpc(FN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = Time('2021-01-01 00:00')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, True)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# jit warmup & test\n",
    "\n",
    "_a = mpc_orbits[0].propagate(t)\n",
    "ar, av = _a.r.to(u.km).value, _a.v.to(u.km / u.s).value\n",
    "_br, _bv = propagate_many(mpc_orbits[:1], t)\n",
    "br, bv = _br.to(u.km).value, _bv.to(u.km / u.s).value\n",
    "\n",
    "np.all(np.isclose(ar, br)), np.all(np.isclose(av, bv))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35.7 s ± 736 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "def propagate_classic(orbits, epoch):\n",
    "    return [orbit.propagate(epoch) for orbit in orbits]\n",
    "\n",
    "%timeit propagate_classic(mpc_orbits, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.38 s ± 38.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit propagate_many(mpc_orbits, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20250"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(mpc_orbits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Propagating a single orbit with `propagate_cpu` to 20250 epochs requires 93.9 ms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "36.38320775026911"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "3380 / 92.9"
   ]
  }
 ],
 "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
}