{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Explicit ODE Schemes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Schemes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explicit Euler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def explicit_euler(y0 : np.ndarray, t : float, f, dt : float):\n",
    "    return y0 + f(y0, t) * dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Runge-Kutta 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runge_kutta_2(y0 : np.ndarray, t : float, f, dt : float):\n",
    "    k1 = f(y0, t)\n",
    "    k2 = f(y0 + k1 * dt, t + dt)\n",
    "    return y0 + (k1 + k2)/2 * dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Runge-Kutta 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runge_kutta_4(y0 : np.ndarray, t : float, f, dt : float):\n",
    "    k1 = f(y0, t)\n",
    "    k2 = f(y0 + k1/2 * dt, t + dt/2)\n",
    "    k3 = f(y0 + k2/2 * dt, t + dt/2)\n",
    "    k4 = f(y0 + k3 * dt, t + dt)\n",
    "    return y0 + (k1 + 2*k2 + 2*k3 + k4)/6 * dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Leapfrog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def leapfrog(y0 : np.ndarray, t : float, f, dt : float):\n",
    "    \n",
    "    # split y0 into position and velocity\n",
    "    r_0 = y0[:2]\n",
    "    v_0 = y0[2:]\n",
    "\n",
    "    def f_v(r):\n",
    "        return f(np.concatenate((r, (0, 0))), t)[2:]\n",
    "\n",
    "    # take the first leap for the velocity\n",
    "    v_1_div_2 = v_0 + f_v(r_0) * dt/2\n",
    "    r_1 = r_0 + v_1_div_2 * dt\n",
    "    v_1 = v_1_div_2 + f_v(r_1) * dt/2\n",
    "\n",
    "    return np.concatenate((r_1, v_1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Semi-implicit Euler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def semi_implicit_euler(y0 : np.ndarray, t : float, f, dt : float):\n",
    "    \n",
    "    # split y0 into position and velocity\n",
    "    r_0 = y0[:2]\n",
    "    v_0 = y0[2:]\n",
    "\n",
    "    # split the function into position and velocity parts\n",
    "    f_r = lambda v : f(np.concatenate(((0, 0), v)), t)[:2]\n",
    "    f_v = lambda r : f(np.concatenate((r, (0, 0))), t)[2:]\n",
    "\n",
    "    # take the leaps\n",
    "    v_1 = v_0 + f_v(r_0) * dt\n",
    "    r_1 = r_0 + f_r(v_1) * dt\n",
    "\n",
    "    return np.concatenate((r_1, v_1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The user calls the wrapper function `integrate`, which, in turn, calls the appropriate scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate(y0 : np.ndarray, f, method : str, dt : float, t_end : float):\n",
    "    \"\"\"Integrate the ODE y' = f(y, t) up to t_end with timestep dt using the given method.\"\"\"\n",
    "    \n",
    "    t = np.arange(0, t_end, dt)\n",
    "    y = np.zeros((len(t), len(y0)))\n",
    "    y[0] = y0\n",
    "    \n",
    "    if method == \"explicit_euler\":\n",
    "        integrator = explicit_euler\n",
    "        \n",
    "    elif method == \"runge_kutta_2\":\n",
    "        integrator = runge_kutta_2\n",
    "\n",
    "    elif method == \"runge_kutta_4\":\n",
    "        integrator = runge_kutta_4\n",
    "\n",
    "    elif method == \"leapfrog\":\n",
    "        integrator = leapfrog\n",
    "\n",
    "    elif method == \"semi_implicit_euler\":\n",
    "        integrator = semi_implicit_euler\n",
    "\n",
    "    for i in range(1, len(t)):\n",
    "        y[i] = integrator(y0 = y[i-1], t = t[i], f = f, dt = dt)\n",
    "    \n",
    "    return t, y\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Kepler Problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial conditions and setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Equations of motion: $\\ddot{\\vec{r}} = -\\frac{G M}{r^3} \\vec{r}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# inital conditions\n",
    "# implicit: GM = 1\n",
    "# star is at rest at the origin\n",
    "\n",
    "#planet\n",
    "e = 0.5 # eccentricity\n",
    "r0 = np.array([1, 0]) # planet's position\n",
    "v0 = np.array([0, np.sqrt(1+e)]) # planet's velocity\n",
    "\n",
    "y0 = np.array([r0, v0]).flatten() # the flatten method converts the 2D array into a 1D array\n",
    "\n",
    "def f(y_vec, t):\n",
    "    x, y, vx, vy = y_vec\n",
    "    r_pow_3_div_2 = np.linalg.norm([x, y])**3\n",
    "    return np.array([vx, vy, - x / r_pow_3_div_2, - y / r_pow_3_div_2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\FLOORV~1\\AppData\\Local\\Temp/ipykernel_10040/3945168747.py:16: RuntimeWarning: invalid value encountered in double_scalars\n",
      "  return np.array([vx, vy, - x / r_pow_3_div_2, - y / r_pow_3_div_2])\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 47480 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t_orbit = 2*np.pi\n",
    "dt = t_orbit/1000\n",
    "t_end = 3*t_orbit\n",
    "\n",
    "t_ee, y_ee = integrate(y0 = y0, f = f, method = \"explicit_euler\", dt = dt, t_end = t_end)\n",
    "t_rk2, y_rk2 = integrate(y0, f, \"runge_kutta_2\", dt = dt, t_end = t_end)\n",
    "t_rk4, y_rk4 = integrate(y0, f, \"runge_kutta_4\", dt = dt, t_end = t_end)\n",
    "t_leapfrog, y_leapfrog = integrate(y0, f, \"leapfrog\", dt = dt, t_end = t_end)\n",
    "t_semi_implicit, y_semi_implicit = integrate(y0, f, \"semi_implicit_euler\", dt = dt, t_end = t_end)\n",
    "\n",
    "fig, ax = plt.subplots(figsize = (10,10))\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.plot(y_ee[:,0], y_ee[:,1], label = \"Explicit Euler\")\n",
    "ax.plot(y_rk2[:,0], y_rk2[:,1], label = \"Runge-Kutta 2\")\n",
    "ax.plot(y_rk4[:,0], y_rk4[:,1], label = \"Runge-Kutta 4\")\n",
    "ax.plot(y_leapfrog[:,0], y_leapfrog[:,1], label = \"Leapfrog\")\n",
    "ax.plot(y_semi_implicit[:,0], y_semi_implicit[:,1], label = \"Semi-Implicit Euler\")\n",
    "ax.set_title(\"Orbit of a planet around a star\")\n",
    "ax.set_xlim(-4.2, 1.2)\n",
    "ax.set_ylim(-3, 3)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rebound example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rebound\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Create a simulation\n",
    "sim = rebound.Simulation()\n",
    "\n",
    "# Set simulation time\n",
    "time = 1e3\n",
    "\n",
    "# Set units\n",
    "sim.units = ('yr', 'AU', 'Msun')\n",
    "\n",
    "# Add particles: Sun and Jupiter\n",
    "sim.add(m=1, hash=\"Sun\")\n",
    "sim.add(m=1e-3, a=5.5, r=4.676e-4, hash=\"Jupiter\")\n",
    "\n",
    "# Move to the barycenter frame\n",
    "sim.move_to_com()\n",
    "\n",
    "# Choose a valid integrator\n",
    "sim.integrator = \"ias15\"  # 'ias15' is a general-purpose, high-accuracy integrator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25152 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set up simulation time and time steps\n",
    "time = 1e3  # Total time to simulate (years)\n",
    "N_steps = 1000  # Number of steps for plotting\n",
    "times = np.linspace(0, time, N_steps)\n",
    "\n",
    "# Arrays to store positions\n",
    "x_jupiter = []\n",
    "y_jupiter = []\n",
    "\n",
    "# Integrate while collecting data\n",
    "for t in times:\n",
    "    sim.integrate(t)\n",
    "    jupiter = sim.particles[1]\n",
    "    x_jupiter.append(jupiter.x)\n",
    "    y_jupiter.append(jupiter.y)\n",
    "\n",
    "# Plot results\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.plot(x_jupiter, y_jupiter, label=\"Jupiter's Orbit\", color=\"blue\")\n",
    "plt.scatter(0, 0, color=\"orange\", label=\"Sun\", s=100)  # Sun at the origin\n",
    "plt.xlabel(\"x (AU)\")\n",
    "plt.ylabel(\"y (AU)\")\n",
    "plt.title(\"Jupiter's Orbit Around the Sun\")\n",
    "plt.axis(\"equal\")\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
