workshop.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries and define a class for the objects\n",
    "\n",
    "We will use numpy's vector operations to speed up calculations. Matplotlib's FuncAnimation will allow us to render a GIF of our system as it evolves with time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from mpl_toolkits import mplot3d\n",
    "\n",
    "class cel_obj:\n",
    "    def __init__(obj, x, y, z, v_x, v_y, v_z, m):\n",
    "        obj.pos = np.array([x, y, z])\n",
    "        obj.vel = np.array([v_x, v_y, v_z])\n",
    "        obj.m = m"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize time variables and G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_0 = 0 \n",
    "t = t_0\n",
    "div = 1\n",
    "dt = 86400/div\n",
    "t_end = 86400 * 365 * 1\n",
    "t_array = np.arange(t_0, t_end, dt)\n",
    "BIG_G = 6.67e-11"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the initial state of the planets\n",
    "\n",
    "Data has been pulled from [NSSDCA](https://nssdc.gsfc.nasa.gov/planetary/factsheet/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sun = cel_obj(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.989e30)\n",
    "mercury = cel_obj(57.909e9, 0.0, 0.0, 0.0, 47.36e3, 0.0, 0.33011e24)\n",
    "venus = cel_obj(108.209e9, 0.0, 0.0, 0.0, 35.02e3, 0.0, 4.8675e24)\n",
    "earth = cel_obj(149.596e9, 0.0, 0.0, 0.0, 29.78e3, 0.0, 5.9724e24)\n",
    "mars = cel_obj(227.923e9, 0.0, 0.0, 0.0, 24.07e3, 0.0, 0.64171e24)\n",
    "jupiter = cel_obj(778.570e9, 0.0, 0.0, 0.0, 13e3, 0.0, 1898.19e24)\n",
    "saturn = cel_obj(1433.529e9, 0.0, 0.0, 0.0, 9.68e3, 0.0, 568.34e24)\n",
    "uranus = cel_obj(2872.463e9, 0.0, 0.0, 0.0, 6.80e3, 0.0, 86.813e24)\n",
    "neptune = cel_obj(4495.060e9, 0.0, 0.0, 0.0, 5.43e3, 0.0, 102.413e24)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create arrays to store the timeseries data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orbital_entities = [sun, mercury, venus, earth, mars, jupiter, saturn, uranus, neptune]\n",
    "ke = []\n",
    "pe = []\n",
    "x_pos = [[],[],[],[]]\n",
    "y_pos = [[],[],[],[]]\n",
    "z_pos = [[],[],[],[]]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evolve the system!\n",
    "\n",
    "We use a simple Euler integrator for the purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while t<t_end:\n",
    "    k = 0\n",
    "    p = 0\n",
    "    a_g = np.zeros_like(orbital_entities)\n",
    "    for x in orbital_entities:\n",
    "        k += 0.5 * x.m * (np.linalg.norm(x.vel))**2\n",
    "    for i in range(len(orbital_entities)):\n",
    "        for j in range(len(orbital_entities)):\n",
    "            if i!=j:\n",
    "                dist = np.linalg.norm(orbital_entities[i].pos - orbital_entities[j].pos) + 0.0001\n",
    "                a_g[i] += (orbital_entities[i].pos - orbital_entities[j].pos)*(-1 * BIG_G * orbital_entities[j].m)/dist**3\n",
    "                p += ((-1 * BIG_G * orbital_entities[i].m * orbital_entities[j].m)/dist)*0.5\n",
    "    ke.append(k/1e35)\n",
    "    pe.append(p/1e35)\n",
    "    for m1_id in range(len(orbital_entities)):                 \n",
    "        orbital_entities[m1_id].vel += a_g[m1_id] * dt\n",
    "    for e_id in range(len(orbital_entities)):\n",
    "        orbital_entities[e_id].pos += orbital_entities[e_id].vel * dt\n",
    "    for i in range(1, 5):\n",
    "        x_pos[i-1].append(orbital_entities[i].pos[0])\n",
    "        y_pos[i-1].append(orbital_entities[i].pos[1])\n",
    "        z_pos[i-1].append(orbital_entities[i].pos[2])\n",
    "    t += dt\n",
    "e = []\n",
    "for i in range(len(ke)):\n",
    "    e.append(ke[i] + pe[i])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the orbits\n",
    "Don't worry about the ffmpeg warning, just know that the GIF is saved in the directory where this notebook is present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(dpi=600)\n",
    "ax = plt.axes(projection='3d')\n",
    "ax.axes.set_xlim3d(left=-2e11, right=2e11) \n",
    "ax.axes.set_ylim3d(bottom=-2e11, top=2e11) \n",
    "ax.axes.set_zlim3d(bottom=-1, top=1)\n",
    "colors = ['red', 'yellow', 'blue', 'orange']\n",
    "\n",
    "def anim_func(i):\n",
    "    for j in range(4):    \n",
    "        ax.plot3D(x_pos[j][:i], y_pos[j][:i], z_pos[j][:i], color=colors[j])\n",
    "        \n",
    "animation = FuncAnimation(fig, anim_func, interval = 50)\n",
    "animation.save(\"solar_system.gif\", dpi=300)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3 = plt.figure()\n",
    "plt.plot(t_array, pe)\n",
    "plt.plot(t_array, ke)\n",
    "plt.plot(t_array, e)\n",
    "fig4 = plt.figure()\n",
    "plt.plot(t_array, ke)\n",
    "fig5 = plt.figure()\n",
    "plt.plot(t_array, pe)\n",
    "fig6 = plt.figure()\n",
    "plt.plot(t_array, e)"
   ]
  }
 ],
 "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.11.0 (main, Oct 24 2022, 18:26:48) [MSC v.1933 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "cb3a50f5fc0d0b888443d45616ff735f38d905fecb7e43e70ba03e33c5084024"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}