{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Orbiting the Sun"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Modeling and Simulation in Python*\n",
    "\n",
    "Copyright 2021 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# install Pint if necessary\n",
    "\n",
    "try:\n",
    "    import pint\n",
    "except ImportError:\n",
    "    !pip install pint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# download modsim.py if necessary\n",
    "\n",
    "from os.path import basename, exists\n",
    "\n",
    "def download(url):\n",
    "    filename = basename(url)\n",
    "    if not exists(filename):\n",
    "        from urllib.request import urlretrieve\n",
    "        local, _ = urlretrieve(url, filename)\n",
    "        print('Downloaded ' + local)\n",
    "    \n",
    "download('https://github.com/AllenDowney/ModSimPy/raw/master/modsim.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import functions from modsim\n",
    "\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a previous example, we modeled the interaction between the Earth and the Sun, simulating what would happen if the Earth stopped in its orbit and fell straight into the Sun.\n",
    "\n",
    "Now let's extend the model to two dimensions and simulate one revolution of the Earth around the Sun, that is, one year.\n",
    "\n",
    "At perihelion, the distance from the Earth to the Sun is 147.09 million km and its velocity is 30,290 m/s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_0 = 147.09e9     # initial distance m\n",
    "v_0 = 30.29e3      # initial velocity m/s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the other constants we'll need, all with about 4 significant digits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = 6.6743e-11     # gravitational constant N / kg**2 * m**2\n",
    "m1 = 1.989e30      # mass of the Sun kg\n",
    "m2 = 5.972e24      # mass of the Earth kg\n",
    "t_end = 3.154e7    # one year in seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Put the initial conditions in a `State` object with variables `x`, `y`, `vx`, and `vy`.\n",
    "Create a `System` object with variables `init` and `t_end`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Write a function called `universal_gravitation` that takes a `State` and a `System` and returns the gravitational force of the Sun on the Earth as a `Vector`.\n",
    "\n",
    "Test your function with the initial conditions; the result should be a Vector with approximate components:\n",
    "\n",
    "```\n",
    "x   -3.66e+22\n",
    "y   0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Write a slope function that takes a timestamp, a `State`, and a `System` and computes the derivatives of the state variables.\n",
    "\n",
    "Test your function with the initial conditions.  The result should be a sequence of four values, approximately\n",
    "\n",
    "```\n",
    "0.0, -30290.0, -0.006, 0.0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**  Use `run_solve_ivp` to run the simulation.\n",
    "Save the return values in variables called `results` and `details`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the following function to plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import plot\n",
    "\n",
    "def plot_trajectory(results):\n",
    "    x = results.x / 1e9\n",
    "    y = results.y / 1e9\n",
    "\n",
    "    make_series(x, y).plot(label='orbit')\n",
    "    plot(0, 0, 'yo')\n",
    "\n",
    "    decorate(xlabel='x distance (million km)',\n",
    "             ylabel='y distance (million km)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_trajectory(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will probably see that the earth does not end up back where it started, as we expect it to after one year.\n",
    "The following cells compute the error, which is the distance between the initial and final positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "error = results.iloc[-1] - system.init\n",
    "error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "offset = Vector(error.x, error.y)\n",
    "vector_mag(offset) / 1e9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem is that the algorithm used by `run_solve_ivp` does not work very well with systems like this.\n",
    "There are two ways we can improve it.\n",
    "\n",
    "`run_solve_ivp` takes a keyword argument, `rtol`, that specifies the \"relative tolerance\", which determines the size of the time steps in the simulation.  Lower values of `rtol` require smaller steps, which yield more accurate results.\n",
    "The default value of `rtol` is `1e-3`.  \n",
    "\n",
    "**Exercise:** Try running the simulation again with smaller values, like `1e-4` or `1e-5`, and see what effect it has on the magnitude of `offset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The other way to improve the results is to use a different algorithm.  `run_solve_ivp` takes a keyword argument, `method`, that specifies which algorithm it should use.  The default is `RK45`, which is a good, general-purpose algorithm, but not particularly good for this system.  One of the other options is `RK23`, which is usually less accurate than `RK45` (with the same step size), but for this system it turns out to be unreasonably good, [for reasons I don't entirely understand](https://mathoverflow.net/questions/314940/celestial-mechanics-and-runge-kutta-methods).\n",
    "Yet another option is 'DOP853', which is particularly good when `rtol` is small.\n",
    "\n",
    "**Exercise:** Run the simulation with one of these methods and see what effect it has on the results.  To get a sense of how efficient the methods are, display `details.nfev`, which is the number of times `run_solve_ivp` called the slope function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "details.nfev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation\n",
    "\n",
    "You can use the following draw function to animate the results, if you want to see what the orbit looks like (not in real time)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlim = results.x.min(), results.x.max()\n",
    "ylim = results.y.min(), results.y.max()\n",
    "\n",
    "def draw_func(t, state):\n",
    "    x, y, vx, vy = state\n",
    "    plot(x, y, 'b.')\n",
    "    plot(0, 0, 'yo')\n",
    "    decorate(xlabel='x distance (million km)',\n",
    "             ylabel='y distance (million km)',\n",
    "             xlim=xlim,\n",
    "             ylim=ylim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# animate(results, draw_func)"
   ]
  },
  {
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}