{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots, LinearAlgebra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HW 3 Solutions: Numerical Solutions of Classical Equations of Motion\n",
    "### Gabe Schumm\n",
    "### Due October 18th, 2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation Functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "System"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Data structure containing moon info\n",
    "mutable struct Moon \n",
    "    m::Float64 #mass in kg\n",
    "    r::Float64 #radius of orbit in meters\n",
    "    T::Float64 #period in seconds\n",
    "\n",
    "    x::Array{Float64} #position vector\n",
    "    v::Array{Float64} #velocity vector\n",
    "\n",
    "    α::Int64 #angle w.r.t. the equitorial plane in DEGREES\n",
    "\n",
    "end\n",
    " \n",
    "#Generator of Moon object\n",
    "function Moon(α::Int64)\n",
    "    \n",
    "    m = 7.3477e22 \n",
    "    r = 384400e3\n",
    "    T = 27.25 * 86164 \n",
    "    \n",
    "    x =  zeros(Float64, 3) \n",
    "    v =  zeros(Float64, 3) \n",
    "    \n",
    "    Moon(m, r, T, x, v, α)\n",
    "end\n",
    "\n",
    "#Data structure containing satellite info\n",
    "mutable struct Satellite\n",
    "\n",
    "    x::Vector{Float64} #position vector\n",
    "    v::Vector{Float64} #velocity vector\n",
    "    \n",
    "    last_phi::Float64 #store last phi value to determine when to increment n_r\n",
    "    n_r::Int64 #number of revolutions around the earth\n",
    "\n",
    "end\n",
    "\n",
    "#Generator of Satellite object\n",
    "function Satellite()\n",
    "    \n",
    "    x =  zeros(Float64, 3)\n",
    "    v =  zeros(Float64, 3)\n",
    "    \n",
    "    Satellite(x, v, 0, 0)\n",
    "end\n",
    "\n",
    "#Data structure containing info for simulation\n",
    "mutable struct System\n",
    "    G::Float64 #gravity constant\n",
    "    M_e::Float64 #mass of earth\n",
    "    T_s::Int64 #sidereal day in seconds\n",
    "\n",
    "    moon::Moon #Moon object used in simulation\n",
    "    sat::Satellite #Satellite object used in simulation\n",
    "    \n",
    "    a::Float64 #initial orbit parameter\n",
    "    \n",
    "    t::Float64 #instaneous time of simulation\n",
    "    dt::Float64 #time step\n",
    "    t_max::Float64 #total total of simulation\n",
    "\n",
    "    N_w::Int64 #cadence of data writing for simulation\n",
    "\n",
    "    outfile::String #data output filename\n",
    "\n",
    "end\n",
    "\n",
    "\n",
    "\n",
    "\n",
    " \n",
    "\n",
    "#Generator of System object\n",
    "#N_t = fraction of sidereal day used for dt\n",
    "#N_max = number of days to integrate over\n",
    "function System(α, a, N_t, N_max, N_w, outfile=false) \n",
    "    G = 6.6743e-11 # m^3 / kg^2\n",
    "    M_e = 5.9736e24 # kg\n",
    "    T_s = 86164 # seconds\n",
    "    \n",
    "    dt = T_s/N_t\n",
    "    t_max = N_max * T_s\n",
    "    \n",
    "    \n",
    "    moon = Moon(α)\n",
    "    sat = Satellite()\n",
    "    \n",
    "\n",
    "    if !isdir(string(\"data/alpha_\", α))\n",
    "        mkdir(string(\"data/alpha_\", α))\n",
    "    end\n",
    "    \n",
    "    outfile = string(\"data/alpha_\", α,\"/sat_a\", a, \"Nt\", N_t, \".dat\")\n",
    "    \n",
    "    \n",
    "    System(G, M_e, T_s, moon, sat, a, 0, dt, t_max, N_w, outfile)\n",
    "end\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "run_simulation (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Returns acceleration of satellite at current time\n",
    "function get_accel(sys)\n",
    "    moon = sys.moon\n",
    "    sat = sys.sat\n",
    "    \n",
    "    r_s = norm(sat.x)\n",
    "\n",
    "    x_sm = sat.x .- moon.x\n",
    "    r_sm = norm(x_sm)\n",
    "    \n",
    "    return (-(sys.G * sys.M_e) * sat.x / r_s^3) - ((sys.G * moon.m) * x_sm / r_sm^3)  #3-element vector\n",
    "end\n",
    "\n",
    "\n",
    "#Initialize moon position and satellite position.velocity\n",
    "function init_system(sys)\n",
    "    moon, sat = sys.moon, sys.sat\n",
    "    \n",
    "    #set moon position according to provided equation\n",
    "    moon.x = [moon.r * cosd(moon.α), 0, moon.r * sind(moon.α)]\n",
    "    \n",
    "    #r_geo and v_geo according to provided equation\n",
    "    #if a≠1, the intitial orbit is not the geo-stationary orbit\n",
    "    r_geo = ((sys.G * sys.M_e * (sys.a*sys.T_s)^2) / (4 * pi^2)) ^(1/3)\n",
    "    v_geo = sqrt(sys.G * sys.M_e / r_geo)\n",
    "    \n",
    "    #set satellite position \n",
    "    sat.x =  [r_geo, 0, 0]\n",
    "    #calculate satellite acceleration at t=0 \n",
    "    accel_0 = get_accel(sys)\n",
    "    #integrate backwards half a time step to calculate v_(-1/2) for leapfrog algorithm\n",
    "    sat.v = [0,v_geo ,0] .- (0.5 * sys.dt) * accel_0\n",
    "\n",
    "end\n",
    "\n",
    "#Update moon position according to current time\n",
    "function update_moon!(sys)\n",
    "    moon = sys.moon\n",
    "    \n",
    "    t_n = sys.t\n",
    "    \n",
    "    moon.x = [moon.r * cosd(moon.α) * cos(2 * pi * t_n/moon.T),\n",
    "        moon.r * sin(2 * pi * t_n/moon.T),\n",
    "        moon.r * sind(moon.α) * cos(2 * pi * t_n/moon.T)]\n",
    "end\n",
    "\n",
    "#Leapfrog integration algorithm\n",
    "function leapfrog!(sys)\n",
    "    sat = sys.sat\n",
    "    \n",
    "    accel = get_accel(sys) #first calculate acceleration a_n\n",
    "    sat.v .+= sys.dt * accel #advance velocity to get v_(n+1/2)\n",
    "    sat.x .+= sys.dt * sat.v #advance position to get x_n+1\n",
    "end\n",
    "\n",
    "#Update full system one time step\n",
    "function update_sys!(sys)\n",
    "    sys.t += sys.dt #increase system time by dt\n",
    "    update_moon!(sys) #update moon position using equation of motion\n",
    "    leapfrog!(sys) #update satellite position using leapfrog algorithm\n",
    "end\n",
    "\n",
    "#Calculate relevant data and write it to file f\n",
    "function write_data(sys, f)\n",
    "    t = sys.t #column 1\n",
    "    phi = mod(atan(sys.sat.x[2], sys.sat.x[1]), 2*pi) #converts from [-π, π] to [0, 2π], column 2\n",
    "    r = norm(sys.sat.x) #in METERS column 3\n",
    "    \n",
    "    #phi_(n+1) - phi_n will always be positive,\n",
    "    #unless phi_(n+1) is greater than 2π aka in the first qudrant of the x-y plane\n",
    "    #if this difference in negative, then the satellite will have just completed a revolution\n",
    "    #and n_r in incremented by one\n",
    "    if phi - sys.sat.last_phi<0  \n",
    "        sys.sat.n_r += 1\n",
    "    end\n",
    "    \n",
    "    #2π*n_r + phi gives cumulative phi\n",
    "    #2π t/T_s gives expected cumulative phi\n",
    "    delta_phi = (2*pi * sys.sat.n_r) + phi -  (2*pi  * sys.t / sys.T_s) #column 4\n",
    "    \n",
    "    \n",
    "    delta_r = r - ((sys.G * sys.M_e * sys.T_s^2) / (4 * pi^2)) ^(1/3) #column 5\n",
    "    \n",
    "    #polar angle theta in spherical coordinates, measured from z-axis\n",
    "    #π/2 - theta gives deviation from equitorial plane \n",
    "    Theta = pi/2 - mod(acos(sys.sat.x[3]/r), 2*pi) #column 6\n",
    "    \n",
    "    sys.sat.last_phi = phi #set current phi to last_phi for next run\n",
    "    \n",
    "    #write data to provided file f\n",
    "    println(f, t, \" \", phi, \" \", r, \" \", delta_phi, \" \", delta_r, \" \", Theta)\n",
    "  \n",
    "end\n",
    "\n",
    "#Run simulation with provided parameters in sys\n",
    "function run_simulation(sys)\n",
    "    open(sys.outfile, \"w\") do f #open outfile as f\n",
    "        for i = 0:(sys.t_max/sys.dt -1) #loop over number of integration steps necesary to reach prodivded t_max\n",
    "            if mod(i, sys.N_w) == 0 #if integration step is multiple of N_w, write relevant data to outfile f\n",
    "                write_data(sys, f)\n",
    "            end\n",
    "            update_sys!(sys) #update system at each integration step\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## $\\alpha = 0$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### $a=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Vector{Float64}:\n",
       "    0.9658300777510512\n",
       " 3074.9068178570305\n",
       "    0.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys = System(0, 1, 10^4, 100, 10^2); #α, a, N_t for dt, N_max for t_max, N_w for data writing\n",
    "\n",
    "#Begin by initializing system (moon and satellite pos/velo)\n",
    "init_system(sys)\n",
    "#Then run simulation\n",
    "#run_simulation(sys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Vector{Float64}:\n",
       " 4.2167508691982634e7\n",
       " 0.0\n",
       " 0.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys.sat.x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing different $a$ Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: norm not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: norm not defined",
      "",
      "Stacktrace:",
      " [1] get_accel(sys::System)",
      "   @ Main ./In[2]:6",
      " [2] init_system(sys::System)",
      "   @ Main ./In[2]:30",
      " [3] top-level scope",
      "   @ In[4]:2",
      " [4] eval",
      "   @ ./boot.jl:360 [inlined]",
      " [5] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)",
      "   @ Base ./loading.jl:1116"
     ]
    }
   ],
   "source": [
    "sys_a = System(0, 1.000485, 10^4, 500, 10^2);\n",
    "init_system(sys_a)\n",
    "run_simulation(sys_a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## $\\alpha = 25$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### $a=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys = System(25, 1, 10^4, 500, 10^2);\n",
    "init_system(sys)\n",
    "run_simulation(sys);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing $a$ Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys = System(25, .5, 10^4, 500, 10^2);\n",
    "init_system(sys)\n",
    "run_simulation(sys);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing satellite with negative $z$ velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "sys = System(25, 1.000438, 10^4, 100, 10^2);\n",
    "\n",
    "moon, sat = sys.moon, sys.sat\n",
    "\n",
    "moon.x = [moon.r * cosd(moon.α) * cos(2 * pi * sys.t/moon.T),\n",
    "    moon.r * sin(2 * pi * sys.t/moon.T),\n",
    "    moon.r * sind(moon.α) * cos(2 * pi * sys.t/moon.T)]\n",
    "\n",
    "\n",
    "r_geo = ((sys.G * sys.M_e * (sys.a*sys.T_s)^2) / (4 * pi^2)) ^(1/3)\n",
    "v_geo = sqrt(sys.G * sys.M_e / r_geo)\n",
    "\n",
    "sat.x =  [r_geo, 0, 0]\n",
    "accel_0 = get_accel(sys)\n",
    "sat.v = [0,v_geo - v_geo*1e-7 ,-sqrt(v_geo^2-(v_geo - v_geo*1e-7)^2)] .- (0.5 * sys.dt) * accel_0;\n",
    "\n",
    "open(\"data/test_zvel.dat\", \"w\") do f\n",
    "        for i = 0:(sys.t_max/sys.dt -1)\n",
    "            if mod(i, sys.N_w) == 0\n",
    "                write_data(sys, f)\n",
    "            end\n",
    "            update_sys!(sys)\n",
    "        end\n",
    "    end"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.2",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}