Kepler_Two_Body_Problem_3_Invariants.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "130d5cd4",
   "metadata": {},
   "source": [
    "## Example 3. Kepler's Two-Body Problem with Three Invariants"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f8878cf",
   "metadata": {},
   "source": [
    "This notebook solves Kepler's Two-Body problem with three invariants by different RK methods and their relaxation versions:\n",
    "\n",
    "\\begin{align}\n",
    "    \\dot{q_{1}} & = p_{1} \\\\\n",
    "    \\dot{q_{2}} & = p_{2} \\\\\n",
    "    \\dot{p_{1}} & = -\\frac{q_{1}}{\\left(q_{1}^2+q_{2}^2\\right)^{3/2}} \\\\\n",
    "    \\dot{p_{2}} & =-\\frac{q_{2}}{\\left(q_{1}^2+q_{2}^2\\right)^{3/2}} \\;,\n",
    "\\end{align}\n",
    "with an initial condition $ \\left(q_{1}(0),q_{2}(0),p_{1}(0),p_{2}(0)\\right)^{T}=\\left(1-e,0,0,\\sqrt{\\frac{1+e}{1-e}}\\right)^{T}$, where $e=0.5$.\n",
    "\n",
    "Three conserved quantities:\n",
    "\n",
    "\\begin{align}\n",
    "    H(q,p) & = \\frac{1}{2}\\left(p_{1}^2+p_{2}^2\\right)-\\frac{1}{\\sqrt{q_{1}^2+q_{2}^2}} \\  \\\\\n",
    "    L(q,p) & = q_{1}p_{2}-q_{2}p_{1} \\\\\n",
    "    A(q,p) & = \\textrm{$L_2$ norm of Laplace–Runge–Lenz vector} \\;\n",
    "\\end{align}\n",
    "\n",
    " where in the last invariant, the well-known Laplace–Runge–Lenz vector function is defined as \n",
    "  \\begin{align}\n",
    "        \\begin{bmatrix}\n",
    "        p_{1} \\\\\n",
    "        p_{2} \\\\\n",
    "        0\n",
    "        \\end{bmatrix} \\times \\begin{bmatrix}\n",
    "        0 \\\\\n",
    "        0 \\\\\n",
    "        q_{1} p_{2}-q_{2} p_{1}\n",
    "        \\end{bmatrix}-\\frac{1}{\\sqrt{q_{1}^{2}+q_{2}^{2}}} \\begin{bmatrix}\n",
    "        q_{1} \\\\\n",
    "        q_{2} \\\\\n",
    "        0\n",
    "        \\end{bmatrix}\\;.\n",
    "   \\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a2ae5d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Required libraries \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import fsolve\n",
    "from nodepy import rk\n",
    "from IPython.display import clear_output\n",
    "from RKSchemes import ssp22, heun33, ssp33, rk44, fehlberg45, dp75"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de4b3520",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Required functions for Kepler's problem\n",
    "def kepler_f(u):\n",
    "    q1 = u[0]\n",
    "    q2 = u[1]\n",
    "    p1 = u[2]\n",
    "    p2 = u[3]\n",
    "    abs_q = np.sqrt(q1*q1 + q2*q2)   \n",
    "    dq1 = p1\n",
    "    dq2 = p2\n",
    "    dp1 = -q1 / (abs_q*abs_q*abs_q)\n",
    "    dp2 = -q2 / (abs_q*abs_q*abs_q)\n",
    "    return np.array([dq1, dq2, dp1, dp2])\n",
    "\n",
    "def kepler_energy_H(u):\n",
    "    abs2_q = u[0]*u[0] + u[1]*u[1]\n",
    "    abs2_p = u[2]*u[2] + u[3]*u[3]\n",
    "    return 0.5 * abs2_p - 1.0 / np.sqrt(abs2_q)\n",
    "\n",
    "def kepler_angular_momentum_L(u):\n",
    "    q1 = u[0]\n",
    "    q2 = u[1]\n",
    "    p1 = u[2]\n",
    "    p2 = u[3]\n",
    "    return q1*p2 - q2*p1\n",
    "\n",
    "def kepler_LRL(u):\n",
    "    q1 = u[0]\n",
    "    q2 = u[1]\n",
    "    p1 = u[2]\n",
    "    p2 = u[3]\n",
    "    abs_q = np.sqrt(q1*q1 + q2*q2)\n",
    "    A1 = p2*(q1*p2 - q2*p1) - q1/abs_q; A2 = -p1*(q1*p2 - q2*p1) - q2/abs_q\n",
    "    return np.linalg.norm(np.array([A1,A2]))**1\n",
    "\n",
    "def rgam(gammas,u,inc1,inc2,inc3,E1_old,E2_old,E3_old):\n",
    "    gamma1, gamma2, gamma3 = gammas\n",
    "    uprop = u + gamma1*inc1 + gamma2*inc2+ gamma3*inc3  \n",
    "    E1 = kepler_energy_H(uprop)\n",
    "    E2 = kepler_angular_momentum_L(uprop)\n",
    "    E3 = kepler_LRL(uprop)\n",
    "    return np.array([E1-E1_old,E2-E2_old,E3-E3_old])\n",
    "\n",
    "\"\"\"\n",
    "Analytical solution of the Kepler problem, cf.\n",
    "http://mathworld.wolfram.com/KeplersEquation.html\n",
    "and\n",
    "https://matlab-monkey.com/astro/keplerEquation/KeplerEquationPub.html\n",
    "\"\"\"\n",
    "def u_kepler(t):\n",
    "    e = 0.5\n",
    "    u0 = np.array([1.0 - e, 0.0, 0.0, np.sqrt((1+e)/(1-e))])\n",
    "    energy = 0.5 * (u0[2]*u0[2] + u0[3]*u0[3]) - 1.0 / np.sqrt(u0[0]*u0[0] + u0[1]*u0[1])\n",
    "    momentum = u0[0]*u0[3] - u0[1]*u0[2]\n",
    "    \n",
    "    res = fsolve(lambda alpha: alpha - e*np.sin(alpha) - t, 0, xtol=1.e-10)\n",
    "    alpha = res[0]\n",
    "    theta = 2.0 * np.arctan(np.sqrt((1+e)/(1-e)) * np.tan(0.5*alpha))\n",
    "    radius = (1 - e*e) / (1 + e * np.cos(theta))\n",
    "    q0 = radius * np.cos(theta)\n",
    "    q1 = radius * np.sin(theta)\n",
    "    \n",
    "    # using the conservation of angular momentum and energy\n",
    "    abs2p = 2 * (energy + 1.0 / np.sqrt(q0*q0 + q1*q1))\n",
    "    if q1 <= 0:\n",
    "        p0 = -((momentum*q1 + np.sqrt(q0**2 * (-momentum**2 + abs2p*(q0**2 + q1**2)))) / (q0**2 + q1**2))\n",
    "        p1 = (momentum*q0**2 - q1*np.sqrt(q0**2 * (-momentum**2 + abs2p*q0**2 + abs2p*q1**2))) / (q0**3 + q0*q1**2)\n",
    "    else:\n",
    "        p0 = (-(momentum*q1) + np.sqrt(q0**2 * (-momentum**2 + abs2p*(q0**2 + q1**2)))) / (q0**2 + q1**2)\n",
    "        p1 = (momentum*q0**2 + q1 * np.sqrt(q0**2 * (-momentum**2 + abs2p*q0**2 + abs2p*q1**2))) / (q0**3 + q0*q1**2)\n",
    "    return np.array([q0, q1, p0, p1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e2ef83",
   "metadata": {},
   "source": [
    "### Baseline RK methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf10593",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute solution with baseline methods\n",
    "def compute_sol_without_relaxation(Mthdname,rkm, dt, f, T, u0,t0): \n",
    "    tt = np.zeros(1) \n",
    "    t = t0; tt[0] = t\n",
    "    uu = np.zeros((1,np.size(u0))) \n",
    "    uu[0,:] = u0.copy()\n",
    "    \n",
    "    s = len(rkm)\n",
    "    y = np.zeros((s,len(u0))) \n",
    "    F = np.zeros((s,len(u0))) \n",
    "    steps = 0\n",
    "    \n",
    "    while t < T and not np.isclose(t, T):\n",
    "        clear_output(wait=True)\n",
    "        if t + dt > T:\n",
    "            dt = T - t\n",
    "        for i in range(s):\n",
    "            y[i,:] = uu[-1].copy()\n",
    "            for j in range(i):\n",
    "                y[i,:] += dt*rkm.A[i,j]*F[j,:]\n",
    "            F[i,:] = f(y[i,:])\n",
    "        inc = dt*sum([rkm.b[i]*F[i] for i in range(s)])    \n",
    "        unew = uu[-1]+inc; t+= dt\n",
    "        tt = np.append(tt, t)\n",
    "        steps +=1\n",
    "        uu = np.append(uu, np.reshape(unew.copy(), (1,len(unew))), axis=0)  \n",
    "        print(\"Method = Baseline %s: Step number = %d (time = %1.2f)\"%(Mthdname,steps,tt[-1]))\n",
    "    return tt, uu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "851bdf5e",
   "metadata": {},
   "source": [
    "### Relaxation RK methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d6118d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computing solution with multiple relaxation methods\n",
    "def compute_sol_multi_relaxation(Mthdname,rkm, dt, f, T, u0, t0):\n",
    "    tt = np.zeros(1) \n",
    "    t = 0; tt[0] = t\n",
    "    uu = np.zeros((1,np.size(u0))) \n",
    "    uu[0,:] = u0.copy()\n",
    "    \n",
    "    s = len(rkm)\n",
    "    y = np.zeros((s,len(u0))) \n",
    "    F = np.zeros((s,len(u0))) \n",
    "    \n",
    "    G1 = np.array([]); G2 = np.array([]) ; G3 = np.array([]) \n",
    "    no_inv = 3; gammas0 = np.zeros(no_inv)\n",
    "    \n",
    "    errs = 0; steps = 0\n",
    "    E10 = kepler_energy_H(u0); E20 = kepler_angular_momentum_L(u0); E30 = kepler_LRL(u0)\n",
    "    \n",
    "    while t < T and not np.isclose(t, T):\n",
    "        clear_output(wait=True)\n",
    "        if t + dt > T:\n",
    "            dt = T - t\n",
    "        for i in range(s):\n",
    "            y[i,:] = uu[-1].copy()\n",
    "            for j in range(i):\n",
    "                y[i,:] += dt*rkm.A[i,j]*F[j,:]\n",
    "            F[i,:] = f(y[i,:])\n",
    "            \n",
    "        inc1 = dt*sum([rkm.b[i]*F[i] for i in range(s)])\n",
    "        inc2 = dt*sum([rkm.bhat[i]*F[i] for i in range(s)])\n",
    "        inc3 = dt*sum([rkm.b3[i]*F[i] for i in range(s)])\n",
    "        \n",
    "        wr_unew = uu[-1] + inc1\n",
    "        E1_old = kepler_energy_H(uu[-1]); E2_old = kepler_angular_momentum_L(uu[-1]); E3_old = kepler_LRL(uu[-1])\n",
    "        \n",
    "        gammas, info, ier, mesg = fsolve(rgam,gammas0,args=(wr_unew,inc1,inc2,inc3,E1_old,E2_old,E3_old),full_output=True)\n",
    "        gamma1, gamma2, gamma3 = gammas\n",
    "             \n",
    "        unew =  wr_unew + gamma1*inc1 + gamma2*inc2+ gamma3*inc3; t+=(1+gamma1+gamma2+gamma3)*dt\n",
    "        tt = np.append(tt, t); steps += 1\n",
    "        uu = np.append(uu, np.reshape(unew.copy(), (1,len(unew))), axis=0) \n",
    "        G1 = np.append(G1, gamma1); G2 = np.append(G2, gamma2); G3 = np.append(G3, gamma3)\n",
    "        print(\"Method = Relaxation %s: At step number = %d (time = %1.2f), integer flag = %d and γ1+γ2+γ3 = %f \\n\"%(Mthdname,steps,tt[-1],ier,gamma1+gamma2+gamma3))\n",
    "    return tt, uu, G1, G2, G3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefd95e8",
   "metadata": {},
   "source": [
    "### Compute solution by all methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "243969c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computing solutions by all the methods\n",
    "%time \n",
    "eqn = 'Kep'\n",
    "methods = [ssp33,fehlberg45,dp75]\n",
    "method_labels = [\"SSPRK(3,3)\", \"Fehlberg(6,4)\", \"DP5(7,5)\"]\n",
    "method_names = [\"SSPRKs3p3\", \"Fehlbergs6p4\", \"DP5s7p5\"]\n",
    "\n",
    "# Inputs to solve the system of ODEs\n",
    "DT = [.05, .05, .1];  f = kepler_f; T = 200;\n",
    "\n",
    "# Initial condition\n",
    "t0 = 0; \n",
    "e = 0.5; u0 = np.array([1.0 - e, 0.0, 0.0, np.sqrt((1+e)/(1-e))])\n",
    "\n",
    "b_tt = []; b_uu = []; r_tt = []; r_uu = []  # empty list to store data by all the methods\n",
    "\n",
    "# Relaxation methods\n",
    "for idx in range(len(methods)):\n",
    "    print(idx)\n",
    "    rkm = methods[idx]; dt = DT[idx]\n",
    "    tt, uu, G1, G2, G3 = compute_sol_multi_relaxation(method_labels[idx], rkm, dt, f, T, u0, t0)\n",
    "    r_tt.append(tt); r_uu.append(uu) \n",
    "    \n",
    "# Baseline methods\n",
    "for idx in range(len(methods)):\n",
    "    rkm = methods[idx]; dt = DT[idx]\n",
    "    tt, uu = compute_sol_without_relaxation(method_labels[idx], rkm, dt, f, T, u0, t0)\n",
    "    b_tt.append(tt); b_uu.append(uu)\n",
    "    \n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b447544c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "path = '%s'%('Figures')\n",
    "\n",
    "import os\n",
    "if not os.path.exists(path):\n",
    "   os.makedirs(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc56a516",
   "metadata": {},
   "source": [
    "### Compute and plot change in invariants by different methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c91da92",
   "metadata": {},
   "outputs": [],
   "source": [
    "b_INV_H = []; b_INV_L = []; b_INV_LRL = []; \n",
    "r_INV_H = []; r_INV_L = []; r_INV_LRL = []; \n",
    "for i in range(len(methods)):\n",
    "    b_inv_H = [kepler_energy_H(u) for u in b_uu[i]] - kepler_energy_H(b_uu[i][0])\n",
    "    b_inv_L = [kepler_angular_momentum_L(u) for u in b_uu[i]] - kepler_angular_momentum_L(b_uu[i][0])\n",
    "    b_inv_LRL = [kepler_LRL(u) for u in b_uu[i]] - kepler_LRL(b_uu[i][0])\n",
    "    r_inv_H = [kepler_energy_H(u) for u in r_uu[i]] - kepler_energy_H(r_uu[i][0])\n",
    "    r_inv_L = [kepler_angular_momentum_L(u) for u in r_uu[i]] - kepler_angular_momentum_L(r_uu[i][0])\n",
    "    r_inv_LRL = [kepler_LRL(u) for u in r_uu[i]] - kepler_LRL(r_uu[i][0])\n",
    "    b_INV_H.append(b_inv_H); b_INV_L.append(b_inv_L); b_INV_LRL.append(b_inv_LRL)\n",
    "    r_INV_H.append(r_inv_H); r_INV_L.append(r_inv_L); r_INV_LRL.append(r_inv_LRL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e8c4e9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "font = {#'family' : 'normal',\n",
    "'weight' : 'normal',\n",
    "'size'   : 14}\n",
    "plt.rc('font', **font) \n",
    "plt.figure(figsize=(15, 4))\n",
    "for i in range(len(methods)):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.plot(b_tt[i],b_INV_H[i],':r',label=\"Baseline: $H(p(t),q(t))-H(p(0),q(0))$\")\n",
    "    plt.plot(r_tt[i],r_INV_H[i],'-r',label=\"Relaxation: $H(p(t),q(t))-H(p(0),q(0))$\")\n",
    "    plt.plot(b_tt[i],b_INV_L[i],':b',label=\"Baseline: $L(p(t),q(t))-L(p(0),q(0))$\")\n",
    "    plt.plot(r_tt[i],r_INV_L[i],'-b',label=\"Relaxation: $L(p(t),q(t))-L(p(0),q(0))$\")\n",
    "    plt.plot(b_tt[i],b_INV_LRL[i],':',color='lime',label=\"Baseline: $A(p(t),q(t))-A(p(0),q(0))$\")\n",
    "    plt.plot(r_tt[i],r_INV_LRL[i],'-',color='lime',label=\"Relaxation:$A(p(t),q(t))-A(p(0),q(0))$\")\n",
    "    plt.title(\"%s with $\\Delta t$ = %.2f\"%(method_labels[i],DT[i]))\n",
    "    plt.xlabel(\"$t$\");\n",
    "    plt.yscale(\"symlog\", linthresh=1.e-14)\n",
    "    plt.yticks([-1.e-6, -1.e-10, -1.e-14, 1.e-14, 1.e-10, 1.e-6])\n",
    "    plt.tight_layout()\n",
    "    \n",
    "ax = plt.gca()\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "plt.figlegend(handles, labels, loc='upper center', ncol=3, bbox_to_anchor=(0.5, 1.2))\n",
    "plt.savefig(\"./Figures/Kepler_Two_Body_3_Inv_Error_Time.pdf\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8335ca25",
   "metadata": {},
   "source": [
    "### Compute and plot errors by different methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76adea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "b_ERR = []; r_ERR = []\n",
    "for i in range(len(methods)):\n",
    "    # maximum norm\n",
    "    b_err = [np.max(np.abs(b_uu[i][j,0:2] - u_kepler(b_tt[i][j])[0:2])) for j in np.arange(len(b_tt[i]))]\n",
    "    r_err = [np.max(np.abs(r_uu[i][j,0:2] - u_kepler(r_tt[i][j])[0:2])) for j in np.arange(len(r_tt[i]))]\n",
    "    b_ERR.append(b_err); r_ERR.append(r_err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ce82df",
   "metadata": {},
   "outputs": [],
   "source": [
    "lgd_box_pos = [[0.4,0.2,0.5, 0.5],[0.4,0.1,0.5, 0.5],[0.4,0.5,0.5, 0.5]]\n",
    "sl1_cons_mult = [8e-6,1e-6,5e-6]; sl1_p = [1,1,1]\n",
    "sl2_cons_mult = [2*1e-4,6e-8,4*1e-7]; sl2_p = [2,2,2]\n",
    "y_scale_line = [1e-10,1e-8,1e-7]\n",
    "shift = 1\n",
    "\n",
    "font = {#'family' : 'normal',\n",
    "'weight' : 'normal',\n",
    "'size'   : 14}\n",
    "plt.rc('font', **font) \n",
    "plt.figure(figsize=(15, 4))\n",
    "\n",
    "for i in range(len(methods)):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.plot(b_tt[i]+shift,b_ERR[i],':',color='orangered',label=\"Baseline\")\n",
    "    plt.plot(r_tt[i]+shift,r_ERR[i],'-g',label=\"Relaxation\")\n",
    "    sl_t = np.linspace(10,200,1000)\n",
    "\n",
    "    plt.plot(sl_t,sl1_cons_mult[i]*sl_t**sl1_p[i],'--',color='0.5',label=\"$\\mathcal{O}(t^{%d})$\"%(sl1_p[i]))\n",
    "    plt.plot(sl_t,sl2_cons_mult[i]*sl_t**sl2_p[i],'-',color='0.5',label=\"$\\mathcal{O}(t^{%d})$\"%(sl2_p[i]))\n",
    "    \n",
    "    plt.xscale(\"log\"); plt.yscale(\"log\")\n",
    "    plt.xlabel('t'); plt.ylabel('Error in q')\n",
    "    plt.title(\"%s with $\\Delta t$ = %.2f\"%(method_labels[i],DT[i]))\n",
    "    plt.tight_layout()\n",
    "    \n",
    "ax = plt.gca()\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "plt.figlegend(handles, labels, loc='upper center', ncol=6, bbox_to_anchor=(0.5, 1.1))\n",
    "plt.savefig(\"./Figures/Kepler_Two_Body_Err_Time.pdf\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4521f1b8",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}