test_guv.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "caf1b6e8-0316-43d5-bc49-c814e1787440",
   "metadata": {},
   "source": [
    "# Setup for this notebook\n",
    "\n",
    "1) Poliastro:\n",
    "\n",
    "```bash\n",
    "git clone git://github.com/s-m-e/poliastro.git\n",
    "cd poliastro\n",
    "git checkout 879f7ab62d05361aff88575bd060d8ff9f880a14\n",
    "pip install -e .\n",
    "```\n",
    "\n",
    "2) Extra packages\n",
    "\n",
    "```bash\n",
    "pip install joblib orbitalpy psutil\n",
    "```\n",
    "\n",
    "3) Data: \n",
    "\n",
    "```bash\n",
    "wget https://minorplanetcenter.net/Extended_Files/nea_extended.json.gz\n",
    "gzip -d nea_extended.json.gz\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53c2ea0c-8aac-438e-957f-9dcc71d2543d",
   "metadata": {},
   "source": [
    "# Loading test data: Near Earth Asteroid (NEA) orbits from the MPC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7eba6c6e-4d4c-49c4-8b03-5b3965f26d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "FN = 'nea_extended.json'\n",
    "\n",
    "import json\n",
    "\n",
    "from astropy import units as u\n",
    "from joblib import Parallel, delayed\n",
    "import numpy as np\n",
    "from orbital.utilities import true_anomaly_from_mean\n",
    "import psutil\n",
    "\n",
    "from poliastro.bodies import Sun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e78db0cc-e163-49a4-88f8-b20f7166b80e",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = Sun.k.to_value(u.km**3 / u.s**2)\n",
    "\n",
    "def _orbit_from_mpc(body):\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",
    "    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(nu.unit)\n",
    "    return (\n",
    "        ((body['a'] * u.AU).to_value(u.km) * (1 - body['e']**2)),  # a\n",
    "        body['e'],  # ecc\n",
    "        (body['i'] * u.deg).to_value(u.rad),  # inc\n",
    "        (body['Node'] * u.deg).to_value(u.rad),  # raan\n",
    "        (body['Peri'] * u.deg).to_value(u.rad),  # argp\n",
    "        nu.to_value(u.rad),  # nu\n",
    "    )\n",
    "\n",
    "def _read_mpc(fn):\n",
    "    with open(fn, 'r', encoding = 'utf-8') as f:\n",
    "        raw = json.load(f)\n",
    "    return Parallel(n_jobs = psutil.cpu_count(logical = True))(delayed(_orbit_from_mpc)(body) for body in raw)\n",
    "\n",
    "mpc_orbits = np.array(_read_mpc(FN), dtype = 'f8').T.copy()  # HACK copy array to force C order"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a433a663-bf20-427f-b585-ed85d21ec7dd",
   "metadata": {},
   "source": [
    "# Import of poliastro's new `jit` infrastructure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "02f18bea-7d02-4519-a937-06e84010894b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu never\n"
     ]
    }
   ],
   "source": [
    "from math import cos, sin, sqrt\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "# POLIASTRO_TARGET can be set to `cpu`, `parallel` or `cuda`. Default: `cpu`\n",
    "os.environ['POLIASTRO_TARGET'] = 'cpu'  \n",
    "\n",
    "# POLIASTRO_INLINE can be set to `always` or `never`. Default: `never`\n",
    "os.environ['POLIASTRO_INLINE'] = 'never'\n",
    "\n",
    "from poliastro.core.jit import gjit, hjit, vjit, TARGET, INLINE\n",
    "\n",
    "print(TARGET, INLINE)  # verfication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "488cd252-9394-4284-b469-74d793d41282",
   "metadata": {},
   "source": [
    "# Isolated code-path for conversion of classical orbital elements to state vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "36393c8e-3aa5-4938-9ba4-3807dff2bc80",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hjit('M(M)')\n",
    "def transpose_M_(a):\n",
    "    return (\n",
    "        (a[0][0], a[1][0], a[2][0]),\n",
    "        (a[0][1], a[1][1], a[2][1]),\n",
    "        (a[0][2], a[1][2], a[2][2]),\n",
    "    )\n",
    "\n",
    "@hjit('M(M,M)')\n",
    "def matmul_MM_(a, b):\n",
    "    return (\n",
    "        (\n",
    "            a[0][0] * b[0][0] + a[0][1] * b[1][0] + a[0][2] * b[2][0],\n",
    "            a[0][0] * b[0][1] + a[0][1] * b[1][1] + a[0][2] * b[2][1],\n",
    "            a[0][0] * b[0][2] + a[0][1] * b[1][2] + a[0][2] * b[2][2],\n",
    "        ),\n",
    "        (\n",
    "            a[1][0] * b[0][0] + a[1][1] * b[1][0] + a[1][2] * b[2][0],\n",
    "            a[1][0] * b[0][1] + a[1][1] * b[1][1] + a[1][2] * b[2][1],\n",
    "            a[1][0] * b[0][2] + a[1][1] * b[1][2] + a[1][2] * b[2][2],\n",
    "        ),\n",
    "        (\n",
    "            a[2][0] * b[0][0] + a[2][1] * b[1][0] + a[2][2] * b[2][0],\n",
    "            a[2][0] * b[0][1] + a[2][1] * b[1][1] + a[2][2] * b[2][1],\n",
    "            a[2][0] * b[0][2] + a[2][1] * b[1][2] + a[2][2] * b[2][2],\n",
    "        ),\n",
    "    )\n",
    "\n",
    "@hjit('V(V,M)')\n",
    "def matmul_VM_(a, b):\n",
    "    return (\n",
    "        a[0] * b[0][0] + a[1] * b[1][0] + a[2] * b[2][0],\n",
    "        a[0] * b[0][1] + a[1] * b[1][1] + a[2] * b[2][1],\n",
    "        a[0] * b[0][2] + a[1] * b[1][2] + a[2] * b[2][2],\n",
    "    )\n",
    "\n",
    "@hjit('M(f,u1)')\n",
    "def rotation_matrix_(angle, axis):\n",
    "    c = cos(angle)\n",
    "    s = sin(angle)\n",
    "    if axis == 0:\n",
    "        return (\n",
    "            (1.0, 0.0, 0.0),\n",
    "            (0.0,   c,  -s),\n",
    "            (0.0,   s,   c),\n",
    "        )\n",
    "    if axis == 1:\n",
    "        return (\n",
    "            (  c, 0.0,   s),\n",
    "            (0.0, 1.0, 0.0),\n",
    "            (  s, 0.0,   c),\n",
    "        )\n",
    "    if axis == 2:\n",
    "        return (\n",
    "            (  c,  -s, 0.0),\n",
    "            (  s,   c, 0.0),\n",
    "            (0.0, 0.0, 1.0),\n",
    "        )\n",
    "    raise ValueError(\"Invalid axis: must be one of 0, 1 or 2\")\n",
    "\n",
    "@hjit('Tuple([V,V])(f,f,f,f)')\n",
    "def rv_pqw_(k, p, ecc, nu):\n",
    "    sinnu = sin(nu)\n",
    "    cosnu = cos(nu)\n",
    "    a = p / (1 + ecc * cosnu)\n",
    "    b = sqrt(k / p)\n",
    "    return (\n",
    "        (cosnu * a, sinnu * a, 0),\n",
    "        (-sinnu * b, (ecc + cosnu) * b, 0),\n",
    "    )\n",
    "\n",
    "@hjit('M(f,f,f)')\n",
    "def coe_rotation_matrix_(inc, raan, argp):\n",
    "    r = rotation_matrix_(raan, 2)\n",
    "    r = matmul_MM_(r, rotation_matrix_(inc, 0))\n",
    "    return matmul_MM_(r, rotation_matrix_(argp, 2))\n",
    "\n",
    "@hjit('Tuple([V,V])(f,f,f,f,f,f,f)')\n",
    "def coe2rv_(k, p, ecc, inc, raan, argp, nu):\n",
    "    \"Converts from classical orbital to state vectors\"\n",
    "    r, v = rv_pqw_(k, p, ecc, nu)\n",
    "    rm = transpose_M_(coe_rotation_matrix_(inc, raan, argp))\n",
    "    return matmul_VM_(r, rm), matmul_VM_(v, rm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cba035a-cf6f-42bf-aabb-d88b981e6b12",
   "metadata": {},
   "source": [
    "# Testing for single orbit/state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "de6b0437-e58a-4aab-b04d-52bcb5e93a64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "((-82217232.08344641, -220850539.44815192, -36798614.1385071), (18.787582906265396, -12.196739597201557, 1.6538450388075674))\n"
     ]
    }
   ],
   "source": [
    "if TARGET != 'cuda':\n",
    "    print(coe2rv_(K, *mpc_orbits[:, 0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ee9e06c-62ff-4778-963f-b25c5c5699aa",
   "metadata": {},
   "source": [
    "# Testing for **array** of orbits/states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6e523d16-a08e-42d0-93b6-de6e820d3d2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float64 True 2 (6, 20250)\n",
      "float64 True 2 (3, 20250)\n",
      "float64 True 2 (3, 20250)\n",
      "float64 True 0 ()\n"
     ]
    }
   ],
   "source": [
    "r = np.zeros((3, mpc_orbits.shape[1]), dtype = 'f8')\n",
    "v = np.zeros((3, mpc_orbits.shape[1]), dtype = 'f8')\n",
    "\n",
    "for item in (mpc_orbits, r, v, K):\n",
    "    print(item.dtype, item.flags.c_contiguous, item.ndim, item.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dc561dec-991b-4e77-ad78-6727ccd31d20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[-8.22172321e+07, -1.68505479e+08,  2.28586647e+08, ...,\n",
       "          3.89940003e+07, -1.25327990e+08,  1.00536243e+08],\n",
       "        [-2.20850539e+08,  3.97092826e+08, -5.12990316e+08, ...,\n",
       "         -1.29757595e+08, -1.72807704e+08,  3.29143864e+08],\n",
       "        [-3.67986141e+07, -8.34073051e+07, -5.81423124e+06, ...,\n",
       "          5.41316389e+07, -4.90166215e+07,  1.21615287e+08]]),\n",
       " array([[ 18.78758291, -16.1191056 ,   8.01970316, ...,  30.09556541,\n",
       "          17.78412559, -16.22194191],\n",
       "        [-12.1967396 ,   2.73889356,   6.87396875, ...,   5.42495671,\n",
       "          -3.92195834,  -5.3500318 ],\n",
       "        [  1.65384504,  -0.78254845,  -1.63941873, ...,   5.21728375,\n",
       "          -1.4502226 ,   4.5105067 ]]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@gjit(\n",
    "    'void(f,f[:],f[:],f[:],f[:],f[:],f[:],f[:],f[:],f[:],f[:],f[:],f[:])',\n",
    "    '(),(n),(n),(n),(n),(n),(n)->(n),(n),(n),(n),(n),(n)',\n",
    ")\n",
    "def coe2rv(k, p, ecc, inc, raan, argp, nu, rx, ry, rz, vx, vy, vz):\n",
    "    \"Converts from classical orbital elements to state vectors ON ARRAYS\"\n",
    "    for idx in range(p.shape[0]):\n",
    "        (rx[idx], ry[idx], rz[idx]), (vx[idx], vy[idx], vz[idx]) = coe2rv_(k, p[idx], ecc[idx], inc[idx], raan[idx], argp[idx], nu[idx])\n",
    "\n",
    "coe2rv(K, *mpc_orbits, *r, *v)\n",
    "r, v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6424d8ac-94d2-4019-b8f5-72a271466f76",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-8.22172321e+07 -2.20850539e+08 -3.67986141e+07] [ 18.78758291 -12.1967396    1.65384504]\n"
     ]
    }
   ],
   "source": [
    "print(r[:,0], v[:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75721bd7-2f8a-47a4-bd09-07baeec6c6f2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}