Eric.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ab20586c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Apr 21 2021 20:26:04\n",
      "WHFast\n",
      "Nplanet:  1  No gradient:  0.0025816455000153837  gradient:  0.009387416999999232  ratio:  3.6362145770762457\n",
      "Nplanet:  2  No gradient:  0.003754166499987832  gradient:  0.02273199999999065  ratio:  6.055139003574915\n",
      "Nplanet:  3  No gradient:  0.0025495829999897524  gradient:  0.037941291999999294  ratio:  14.881371581216142\n",
      "IAS15\n",
      "Nplanet:  1  No gradient:  0.047735312500009286  gradient:  0.2833853959999999  ratio:  5.936598739139809\n",
      "Nplanet:  2  No gradient:  0.0653132080000205  gradient:  0.8796171875000027  ratio:  13.467676974307045\n",
      "Nplanet:  3  No gradient:  0.08157847949999564  gradient:  1.8487465414999917  ratio:  22.662184350960974\n"
     ]
    }
   ],
   "source": [
    "# Compares the run time of REBOUND with gradient\n",
    "# versus NbodyGradient\n",
    "\n",
    "import rebound\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "from timeit import Timer\n",
    "print(rebound.__build__)\n",
    "\n",
    "#def evaluate(order=0,N=0, nplanet = 0, integrator=\"ias15\"):\n",
    "def evaluate(order=0,N=0, nplanet = 0, integrator=\"whfast\"):\n",
    "    sim = rebound.Simulation()\n",
    "    sim.integrator = integrator\n",
    "    sim.ri_whfast.safe_mode = 0\n",
    "    sim.add(m=1.)\n",
    "    semi = 1.\n",
    "    f0 = 0.0\n",
    "    for i in range(nplanet):\n",
    "      sim.add(primary=sim.particles[0],m=1e-5, a=semi, f=f0)\n",
    "      semi *= 1.8\n",
    "      f0 += 1.4\n",
    "    \n",
    "    vs = [\"a\",\"e\",\"i\",\"omega\",\"Omega\",\"f\",\"m\"]\n",
    "    var = []\n",
    "    if order>=1:\n",
    "        for i in range(N):\n",
    "            pi = 1\n",
    "            shifti = 0\n",
    "            if i>=7:\n",
    "                pi = 2\n",
    "                shifti = -7\n",
    "            if i>=14:\n",
    "                pi = 3\n",
    "                shifti = -14\n",
    "            if i>=21:\n",
    "                pi = 4\n",
    "                shifti = -21\n",
    "            if i>=28:\n",
    "                pi = 4\n",
    "                shifti = -28\n",
    "            if i>=35:\n",
    "                pi = 5\n",
    "                shifti = -35\n",
    "            if i>=42:\n",
    "                pi = 6\n",
    "                shifti = -42\n",
    "            if i>=49:\n",
    "                pi = 7\n",
    "                shifti = -49\n",
    "            if i>=56:\n",
    "                pi = 8\n",
    "                shifti = -56\n",
    "            if i>=63:\n",
    "                pi = 9\n",
    "                shifti = -63\n",
    "            if i>=70:\n",
    "                pi = 10\n",
    "                shifti = -70\n",
    "            var_d = sim.add_variation()\n",
    "            var_d.vary(pi,vs[i+shifti])\n",
    "            var.append(var_d)\n",
    "    sim.move_to_com()\n",
    "    sim.dt = np.pi/10.\n",
    "    sim.integrate(5000.)\n",
    "#    if order==0:\n",
    "#        torb = 2.*np.pi\n",
    "#        Noutputs = 100\n",
    "#        times = np.linspace(500.+torb, 500.+2.*torb, Noutputs)\n",
    "#        x = np.zeros([Noutputs,3])\n",
    "#        for i,time in enumerate(times):\n",
    "#           sim.integrate(time,exact_finish_time=1)\n",
    "#           x[i,0] = times[i]\n",
    "#           x[i,1] = sim.particles[1].x    \n",
    "#           x[i,2] = sim.particles[1].y\n",
    "#        np.savetxt('rebound_4body_xy.txt',x,delimiter=',')\n",
    "    return \n",
    "\n",
    "#def evaluateWithN(order,N,integrator=\"ias15\",nplanet = 0):\n",
    "def evaluateWithN(order,N,integrator=\"whfast\",nplanet = 0):\n",
    "    def _e():\n",
    "        evaluate(order,N,nplanet,integrator)\n",
    "        pass\n",
    "    return _e\n",
    "\n",
    "print(\"WHFast\")\n",
    "for i in range(3):\n",
    "  Nmax = (i+1)*7+1\n",
    "  repeat = 1\n",
    "  t = Timer(evaluateWithN(0,0,nplanet=i+1))\n",
    "  var_0 = t.timeit(repeat)/2.\n",
    "  t = Timer(evaluateWithN(1,Nmax-1,nplanet=i+1))\n",
    "  var_1 =  t.timeit(repeat)/2.\n",
    "\n",
    "  print(\"Nplanet: \",i+1,\" No gradient: \",var_0, \" gradient: \", var_1, \" ratio: \", var_1/var_0)\n",
    "\n",
    "print(\"IAS15\")\n",
    "for i in range(3):\n",
    "  Nmax = (i+1)*7+1\n",
    "  repeat = 1\n",
    "  t = Timer(evaluateWithN(0,0,nplanet=i+1,integrator=\"ias15\"))\n",
    "  var_0 = t.timeit(repeat)/2.\n",
    "  t = Timer(evaluateWithN(1,Nmax-1,nplanet=i+1,integrator=\"ias15\"))\n",
    "  var_1 =  t.timeit(repeat)/2.\n",
    "\n",
    "  print(\"Nplanet: \",i+1,\" No gradient: \",var_0, \" gradient: \", var_1, \" ratio: \", var_1/var_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "691b32d1",
   "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": 5
}

Thanks much! @langfzac