Vector Elliptical Adaptive RK4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ellipsefns\n",
    "#getfocus, getsemiminor, theoreticalx, theoreticaly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import initdat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta=0.0001 #with dtmin=0.001, accuracy of 1e-6 appears to be too much to ask. try 0.05\n",
    "rad0=50\n",
    "ecc0=0.5\n",
    "theta0= 0. #math.pi/6. #math.pi/6.\n",
    "mass0=1.0\n",
    "initdateqellipse=initdat.InitialDataEqualMassConic(rad0,ecc0,theta0,mass0)\n",
    "print(initdateqellipse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getxyuvtwoellipse_apper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "useperihelion=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyuvaeqellipse=getxyuvtwoellipse_apper.getxyuveqtwoellipses(initdateqellipse,useperihelion)\n",
    "print(xyuvaeqellipse)\n",
    "masses,xvec,avec=xyuvaeqellipse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import RK4implicit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def RK4implicit(h,t,xvec,f): #not a finite difference so no step in y\n",
    "    k1= h*f(t,xvec)\n",
    "    k2=h*f(t+h/2, xvec+k1/2)\n",
    "    k3=h*f(t+h/2,xvec+k2/2)\n",
    "    k4=h*f(t+h,xvec+k3)\n",
    "    return t+h, xvec+1/6.*(k1+2.*k2+2.*k3+k4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import OrbitDiffEqAdaptive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masses, xvec,avec=xyuvaeqellipse\n",
    "ODEeq= OrbitDiffEqAdaptive.OrbitDiffEqAdaptive(masses,xvec,avec,0.0, delta)\n",
    "ODEeq.print2D()\n",
    "ODEeq.print2D()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(xyuvaeqellipse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import timestepAdaptive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numsteps=500000\n",
    "\n",
    "#numsteps=6000*10\n",
    "\n",
    "dt=.01\n",
    "dtmax=1\n",
    "dtmin = 0.001\n",
    "outputevery=1\n",
    "nptimes, npstar1x, npstar2x, npstar1a, npstar2a,nperr,npdt= timestepAdaptive.timestepAdaptive(numsteps,dt,dtmax,dtmin,xyuvaeqellipse,outputevery,delta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "npstar1x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "t = np.linspace(0, 2*math.pi, 100)\n",
    "xax=-200+np.arange(0,400)\n",
    "plt.ylim(-200,200)\n",
    "plt.xlim(-200,200)\n",
    "plt.title(\"Elliptical Orbit, Adaptive RK4, a=\"+str(rad0)+\", e=\"+ str(ecc0)+\", theta=\"+str(round(180*theta0/math.pi))+ \" degrees\")\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,1),ellipsefns.getytheory(theta0,ecc0,rad0,t,1),'b')\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,-1),ellipsefns.getytheory(theta0,ecc0,rad0,t,-1),'g')\n",
    "plt.plot(npstar1x[:,0],npstar1x[:,1],'y')\n",
    "plt.plot(npstar2x[:,0],npstar2x[:,1],'r')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adaptive RK4 orbits run and superficially look correct. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getfocus(a,ecc):\n",
    "    c=a*ecc\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getsemiminor(a,ecc):\n",
    "    c=getfocus(a,ecc)\n",
    "    b=np.sqrt(a**2-c**2)\n",
    "    return b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getxtheory2(theta0,ecc0,rad0,t, sgn):\n",
    "    focus=getfocus(rad0,ecc0)\n",
    "    b=getsemiminor(rad0,ecc0)\n",
    "    xdim=rad0\n",
    "    ydim=b\n",
    "    xtheory=sgn*focus*np.cos(theta0)+xdim*np.cos(t)*np.cos(theta0)-ydim*np.sin(t)*np.sin(theta0)\n",
    "    return xtheory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getytheory2(theta0, ecc0, rad0, t, sgn):\n",
    "    focus = getfocus(rad0,ecc0)\n",
    "    b=getsemiminor(rad0,ecc0)\n",
    "    xdim=rad0\n",
    "    ydim=b\n",
    "    ytheory=sgn*focus*np.sin(theta0)+ydim*np.cos(theta0)*np.sin(t)+xdim*np.sin(theta0)*np.cos(t)\n",
    "    return ytheory\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "t = np.linspace(0, 2*math.pi, 100)\n",
    "xax=np.arange(0,100*len(npstar2x[:,1]))\n",
    "xax=xax-len(xax)/2\n",
    "oneellipse=int(len(npdt))\n",
    "plt.ylim(-200,200)\n",
    "plt.xlim(-200,200)\n",
    "#plt.ylim(-1,1)\n",
    "#plt.xlim(74.999,75)\n",
    "plt.title(\"Elliptical Orbit, RK4, a=\"+str(rad0)+\", e=\"+ str(ecc0)+\", theta=\"+str(round(180*theta0/math.pi))+ \" degrees\")\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,1),ellipsefns.getytheory(theta0,ecc0,rad0,t,1),'b')\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,-1),ellipsefns.getytheory(theta0,ecc0,rad0,t,-1),'g')\n",
    "plt.plot(npstar1x[:oneellipse,0],npstar1x[:oneellipse,1],'y')\n",
    "plt.plot(npstar2x[:oneellipse,0],npstar2x[:oneellipse,1],'r')\n",
    "plt.plot(xax, np.tan(theta0)*xax,'c')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With adaptive step size I get inspirals!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "t = np.linspace(0, 2*math.pi, 100)\n",
    "xax=np.arange(0,100*len(npstar2x[:,1]))\n",
    "xax=xax-len(xax)/2\n",
    "\n",
    "#plt.ylim(-200,200)\n",
    "#plt.xlim(-200,200)\n",
    "plt.ylim(-1,1)\n",
    "plt.xlim(74.999,75)\n",
    "plt.title(\"Elliptical Orbit, RK4, a=\"+str(rad0)+\", e=\"+ str(ecc0)+\", theta=\"+str(round(180*theta0/math.pi))+ \" degrees\")\n",
    "#plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,1),ellipsefns.getytheory(theta0,ecc0,rad0,t,1),'b')\n",
    "#plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,-1),ellipsefns.getytheory(theta0,ecc0,rad0,t,-1),'g')\n",
    "plt.plot(npstar1x[:oneellipse,0],npstar1x[:oneellipse,1],'*y')\n",
    "plt.plot(npstar2x[:oneellipse,0],npstar2x[:oneellipse,1],'*r')\n",
    "plt.plot(xax, np.tan(theta0)*xax,'c')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "t = np.linspace(0, 2*math.pi, 100)\n",
    "xax=np.arange(0,100*len(npstar2x[:,1]))\n",
    "xax=xax-len(xax)/2\n",
    "\n",
    "#plt.ylim(-200,200)\n",
    "#plt.xlim(-200,200)\n",
    "plt.ylim(-10,10)\n",
    "plt.xlim(-25,-20)\n",
    "plt.title(\"Elliptical Orbit, RK4, a=\"+str(rad0)+\", e=\"+ str(ecc0)+\", theta=\"+str(round(180*theta0/math.pi))+ \" degrees\")\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,1),ellipsefns.getytheory(theta0,ecc0,rad0,t,1),'b')\n",
    "plt.plot(ellipsefns.getxtheory(theta0,ecc0,rad0,t,-1),ellipsefns.getytheory(theta0,ecc0,rad0,t,-1),'g')\n",
    "plt.plot(npstar1x[:,0],npstar1x[:,1],'*y')\n",
    "plt.plot(npstar2x[:,0],npstar2x[:,1],'*r')\n",
    "plt.plot(xax, np.tan(theta0)*xax,'c')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nperr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "plt.xlabel(\"timestep\")\n",
    "plt.ylabel(\"Rel err: 1/30 *|x2-x1|*(h'/h)^5\")\n",
    "plt.title(\"Relative error, RK4\")\n",
    "relerrplt,=plt.plot(np.arange(2,oneellipse),nperr[2:oneellipse],'b',label=\" dtmin = \"+ str(dtmin)+\" delta = \"+str(delta))\n",
    "plt.legend(handles=[relerrplt])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This curve should be constant and about equal to delta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "plt.xlabel(\"timestep number\")\n",
    "plt.ylabel(\"timestep size\")\n",
    "plt.title(\"Evolution of timestep with step\")\n",
    "relerrplt,=plt.plot(np.arange(2,oneellipse),npdt[2:oneellipse],'b',label=\" dtmin = \"+ str(dtmin)+\" delta = \"+str(delta))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For appropriately set tolerances, adaptive timestep no longer has feedback issues. Tol=delta=0.0005 or less. The important thing is that this curve is smooth and has no constant regions where the value is pinned to either the maximum or minimum. We also seek reasonable starting values for the timestep. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"relative error\")\n",
    "plt.title(\"horizontal position vs relative error\")\n",
    "plt.plot(npstar1x[2:oneellipse,0], nperr[2:oneellipse],'b')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"y\")\n",
    "plt.ylabel(\"rel error\")\n",
    "plt.title(\"y vs rel error\")\n",
    "plt.plot(npstar1x[2:oneellipse,1],nperr[2:oneellipse])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"r\")\n",
    "plt.ylabel(\"rel error\")\n",
    "plt.title(\"r vs rel error\")\n",
    "rstar1x=np.sqrt(npstar1x[:,0]**2+npstar1x[:,1]**2)\n",
    "plt.plot(rstar1x[2:oneellipse],nperr[2:oneellipse],'*b')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"t\")\n",
    "plt.ylabel(\"r\")\n",
    "plt.title(\"r vs t\")\n",
    "plt.plot(nptimes[2:oneellipse], rstar1x[2:oneellipse], '*b')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thetaellipse = np.arctan2(npstar1x[2:oneellipse,1],npstar1x[2:oneellipse,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"t\")\n",
    "plt.ylabel(\"theta\")\n",
    "plt.title(\"theta vs t\")\n",
    "plt.plot(nptimes[2:oneellipse],thetaellipse,'b')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"theta\")\n",
    "plt.ylabel(\"rel error\")\n",
    "plt.title(\"rel error vs theta\")\n",
    "plt.plot(thetaellipse, nperr[2:oneellipse],'*b')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conservation of angular momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Lstability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deltaL, Lavg=Lstability.Lstability(npstar1x,npstar2x,mass0)\n",
    "deltaL, Lavg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deltaL*Lavg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Orbit conserves angular momentum to a tolerance of deltaL. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conservation of Energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Estability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deltaE, Eavg = Estability.Estability(npstar1x, npstar2x, mass0)\n",
    "deltaE, Eavg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deltaE*Eavg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Orbit conserves energy to a relative error of deltaE. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Symmetry of the two stars orbits with respect to eachother "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import orbitSymmetricTol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "allcancelsx, allcancelsa=orbitSymmetricTol.howSymmetric(npstar1x,npstar2x,npstar1a,npstar2a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "allcancelsx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Orbit is symmetric with respect to x and v up to a tolerance of allcancelsx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "allcancelsa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "npstar1x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Orbit is symmetric with respect to a to up to or better than a tolerance of allcancelsa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing for precision in perihelion and aphelion recovery from knowledge of the initial angle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aeTest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "plt.plot( np.abs(npstar1x[:,1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findPerapHelionThetaIsZero(npstar1x,ecc0,rad0,useperihelion):\n",
    "    maxindex=len(npstar1x[:,1])\n",
    "    apindex=np.argmin(np.abs(npstar1x[10:int(maxindex/2),1]))+1\n",
    "    print(apindex)\n",
    "    delta=math.floor(.2*apindex)\n",
    "    print(delta)\n",
    "    perindex=np.argmin(np.abs(npstar1x[delta:apindex-delta,1]))+delta\n",
    "    print(perindex)\n",
    "    print(apindex/perindex)\n",
    "    if np.abs(apindex/perindex-2.)>0.05: \n",
    "        perindex=np.argmin(np.abs(npstar1x[apindex+delta:2*apindex+delta,1]))+delta+apindex\n",
    "        print(perindex)\n",
    "        apindex,perindex=perindex,apindex\n",
    "        print(apindex,perindex)\n",
    "        print(apindex/perindex)\n",
    "    #apindex3=np.argmin(np.abs(npstar1x[100:apindex2-100,1]))+100\n",
    "    #apindex4=np.argmin(np.abs(npstar1x[100:apindex3-100,1]))+100\n",
    "    #perindex=np.argmin(np.abs(npstar1x[100:apindex2-100,1]))+100\n",
    "    if useperihelion:\n",
    "        apindex,perindex=perindex,apindex\n",
    "    coordap=[npstar1x[apindex,0],npstar1x[apindex,1]]\n",
    "    coordper=[npstar1x[perindex,0],npstar1x[perindex,1]]\n",
    "    ra=np.sqrt(np.sum(np.array(coordap)**2))\n",
    "    rp=np.sqrt(np.sum(np.array(coordper)**2))\n",
    "    e=(ra-rp)/(ra+rp)\n",
    "    a=1./2.*(ra+rp)\n",
    "    deltae=np.abs(e-ecc0)/ecc0\n",
    "    deltaa=np.abs(a-rad0)/rad0\n",
    "    return apindex,perindex,coordap,coordper,ra,rp,e,a,deltae,deltaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if theta0==0:\n",
    "    apindex,perindex,coordap,coordper,ra,rp,e,a,deltae,deltaa=findPerapHelionThetaIsZero(npstar1x,ecc0,rad0,useperihelion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if theta0!=0:\n",
    "    perindex,apindex,coordper,coordap,rp,ra,e,a,deltae,deltaa=aeTest.findPerApHelion(npstar1x,theta0,ecc0,rad0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recovered perihelion radius given by r_per, recovered aphelion radius given by r_ap, a is semimajor axis computed from the orbit, e is eccentricty computed from the orbit, delta a is relative error on the semimajor axis, and delta e is eccenticity relative error. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "apindex,perindex,coordap,coordper,ra,rp,e,a,deltae,deltaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimes[:perindex],npstar1x[:perindex,1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fit t as a function of the eccentric anomoly (angle from the center of the elipse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findEccentricAnomoly(coordap, coordper):\n",
    "    Ecenter=np.array(coordap)+np.array(coordper)\n",
    "    Ecenter/=2.\n",
    "    return Ecenter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ellipsecenter=findEccentricAnomoly(coordap,coordper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ellipsecenter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getEllipsecenter(ecc,rad,theta, useperihelion):\n",
    "    #ind of use perihelion\n",
    "    rp=rad*(1-ecc)\n",
    "    ra=-rad*(1+ecc)\n",
    "    rcent=(rp+ra)/2\n",
    "    #print(rp,ra,rcent)\n",
    "    rcentx=rcent*np.cos(theta)\n",
    "    rcenty=rcent* np.sin(theta)\n",
    "    return [rcentx,rcenty]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(np.array(coordap)+np.array(coordper))/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getEllipsecenter(ecc0,rad0,theta0, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.arctan2(-1.,0)/math.pi*180"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getEang(npstar2x,ecc,rad,theta, useperihelion):\n",
    "    ellipsecenter=getEllipsecenter(ecc,rad,theta, useperihelion)\n",
    "    #print(ellipsecenter)\n",
    "    #print(npstar1x)\n",
    "    #print(npstar1x[:,0])\n",
    "    xfortan=npstar2x[:,0]-ellipsecenter[0]\n",
    "    yfortan=npstar2x[:,1]-ellipsecenter[1]\n",
    "    #print(\"xy\",xfortan,yfortan)\n",
    "    Eang=np.arctan2(yfortan,xfortan)\n",
    "    Eangtheta=2*np.arctan2(np.sqrt(1-ecc)*np.sin(theta/2.),np.sqrt(1+ecc)*np.cos(theta/2.))\n",
    "    Eang-=Eangtheta\n",
    "    #if useperihelion:\n",
    "    #    Eang-=math.pi\n",
    "    print(len(Eang))\n",
    "    for i in np.arange(len(Eang)):\n",
    "        #if xfortan[i]<0:\n",
    "        #    Eang[i]+=math.pi\n",
    "        if Eang[i]<0:\n",
    "            Eang[i]+=2*math.pi\n",
    "    print(xfortan,yfortan,Eang/math.pi*180)\n",
    "    print(\"Eang\",Eang)\n",
    "    return xfortan, yfortan,Eang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.arctan2(-6.69327422*10**-6, -0.02240526)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def xofeccanom(rad,Eang,ecc):\n",
    "    return rad*(np.cos(Eang)+ecc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def yofeccanom(rad, Eang, ecc):\n",
    "    b=getsemiminor(rad,ecc)\n",
    "    return b*np.sin(Eang) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tofEccAnom(npstar2x, rad,ecc,theta,mass, useperihelion):\n",
    "    xE0,yE0,Eang=getEang(npstar2x,ecc,rad,theta, useperihelion)\n",
    "    b=getsemiminor(rad,ecc)\n",
    "    #fitval=rad*b*(Eang-ecc*np.sin(Eang))\n",
    "    sgn=-1\n",
    "    if useperihelion:\n",
    "        sgn*=-1\n",
    "    # added a factor of 2 in front of fitval\n",
    "    fitval=rad**(1.5)*mass**(-.5)*(Eang-sgn*ecc*np.sin(Eang))\n",
    "    fitval*=2\n",
    "    print(Eang, ecc,rad,mass,fitval)\n",
    "    return fitval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def residualtEang(pars,nptimes,npstar2x):\n",
    "    residual=np.abs(nptimes-tofEccAnom(npstar2x,pars['rad'],pars['ecc'],pars['theta'],pars['mass'],pars['useperihelion']))\n",
    "    #print(nptimes, residual)\n",
    "    return residual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endindex=apindex\n",
    "if useperihelion:\n",
    "    endindex=perindex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endindex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xE0,yE0,Eang0=getEang(npstar2x[0:endindex-1,:],ecc0,rad0,theta0, useperihelion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nptimesoneorbit=nptimes[0:endindex-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(Eang0, xE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(Eang0, yE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(xE0,yE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,xE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,yE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,yE0/xE0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.arctan(0./1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yE0[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xE0[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,np.arctan2(yE0,xE0)*180/math.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Eangtest=tofEccAnom(npstar2x[0:endindex-1,:],rad0,ecc0,theta0,masses[0], useperihelion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.max(Eangtest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nptimes[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nptimesoneorbit=nptimes[0:endindex-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tpredoneorbit=Eangtest[0:endindex-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Eangoneorbit=Eang0[0:endindex-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,npstar2x[0:endindex-1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,Eangoneorbit/math.pi*180)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,tpredoneorbit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(nptimesoneorbit,tpredoneorbit, label='times')\n",
    "plt.xlabel(\"T timesteps\")\n",
    "plt.ylabel(\"T computed from angle\")\n",
    "plt.title(\"Times using start parameters\")\n",
    "plt.show()\n",
    "plt.plot(nptimesoneorbit,tpredoneorbit-nptimesoneorbit, label='residual')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(nptimesoneorbit,tpredoneorbit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmfit import Minimizer,Parameters\n",
    "pfittEang = Parameters()\n",
    "pfittEang.add(name='theta', value=theta0, min=theta0-math.pi, max=theta0+math.pi, vary=True)\n",
    "pfittEang.add(name='rad', value=rad0, min=0, vary=True)\n",
    "pfittEang.add(name='ecc', value=ecc0, min=0, max=1.0, vary=True)\n",
    "pfittEang.add(name='mass', value=masses[0], vary=False)\n",
    "pfittEang.add(name='useperihelion', value=True, vary=False)\n",
    "\n",
    "minit = Minimizer(residualtEang, pfittEang, fcn_args=(nptimesoneorbit,npstar2x[:endindex-1,:]))\n",
    "\n",
    "outt = minit.leastsq()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "outt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Semimajor axis fit is good to about .1 percent at an eccentricity of 0.5. rad=49.886512 for delta=0.0001. rad=49.8863857 for delta = 0.0005. ecc = 0.51764747 for delta = 0.0001. ecc = 0.51765166 for delta = 0.0005. so both ecc and rad improve with a smaller delta, but not quickly. The program runs very slowly for large delta and requires 500000 timesteps for delta=0.0001. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfinal=nptimesoneorbit[endindex-2]\n",
    "tfinal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.fft import fft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import blackman\n",
    "wndow=blackman(len(nptimesoneorbit))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ftresiduals=fft(nptimesoneorbit-tpredoneorbit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xf = np.linspace(0.0, 1.0/(2.0*dt), len(nptimesoneorbit)//2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.semilogy(xf[1:len(nptimesoneorbit)//2], 2.0/len(nptimesoneorbit) * np.abs(ftresiduals[1:len(nptimesoneorbit)//2]), '-b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ftr=fft(np.sqrt(npstar1x[:endindex-1,0]**2+npstar1x[:endindex-1,1]**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"frequency\")\n",
    "plt.ylabel(\"Fourier transform of t(E)\")\n",
    "plt.title(\"Fourier transform of time computed from eccentric anomaly\")\n",
    "plt.semilogy(xf[1:len(nptimesoneorbit)//2], 2.0/len(nptimesoneorbit) * np.abs(ftr[1:len(nptimesoneorbit)//2]), '-b')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the fourier transform of the triangle wave."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(\"Fourier transform of residuals of times\")\n",
    "plt.xlabel(\"frequency\")\n",
    "plt.ylabel(\"FT(t(E))\")\n",
    "plt.semilogy(xf[1:50], 2.0/len(nptimesoneorbit) * np.abs(ftresiduals[1:50]), '-b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"bin number\")\n",
    "plt.ylabel(\"FT(t(E))\")\n",
    "plt.title(\"Fourier transform of residuals of times\")\n",
    "plt.semilogy(np.arange(1,30), 2.0/len(nptimesoneorbit[:29]) * np.abs(ftresiduals[1:30]), '*b')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Peak at the second harmonic and the sixth harmonic. The second harmonic seems natural due to the symmetry of the orbit but the sixth harmonic is harder to explain. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implementing the Adaptive RK4 hasn't changed the fit, residuals, or harmonics, sadly. There is still a 2% error in the fit for an eccentricity of 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}