Add python wrapper for SimulationArchive helper C functions; add notebooks

assist/tools.py

@@ -1,8 +1,14 @@
+
 from ._libassist import clibassist
-from ctypes import POINTER, byref, c_int
+from ctypes import POINTER, byref, c_int, c_double
 import rebound
 
 def simulation_convert_to_rebound(sim, ephem, merge_moon=1):
     clibassist.assist_simulation_convert_to_rebound.restype = POINTER(rebound.Simulation)
     sim2 = clibassist.assist_simulation_convert_to_rebound(byref(sim), byref(ephem), c_int(merge_moon))
     return sim2.contents
+
+def assist_create_interpolated_simulation(sa, t):
+    clibassist.assist_create_interpolated_simulation.restype = POINTER(rebound.Simulation)
+    sim = clibassist.assist_create_interpolated_simulation(byref(sa), c_double(t))
+    return sim.contents

jupyter_examples/SimulationArchives.ipynb

@@ -0,0 +1,245 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Saving ASSIST simulation files"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Rebound has the ability to save `SimulationArchive`'s. However, because of the way ASSIST's `integrate_or_interpolate` works (and specifically how it uses data *across* REBOUND steps.\n",
+    "\n",
+    "We will validate this by:\n",
+    "1. doing a traditional ASSIST simulation, `integrate_or_interpolate`ing for a set number of times, and\n",
+    "2. doing a manual step-through simulation, saving it to a file, loading it back up, and calling `integrate_or_interpolate` on the same timesteps,\n",
+    "3. using the C function `assist_create_interpolated_simulation` to use the correct accelerations between steps (see: https://github.com/matthewholman/assist/commit/d94d933357a7dee78347d5e1530d5a641c621bac)\n",
+    "\n",
+    "Then, we'll compare the data."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First, let's import REBOUND, ASSIST as well as numpy and matplotlib."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import rebound\n",
+    "import assist\n",
+    "\n",
+    "ephem = assist.Ephem(\"../data/linux_p1550p2650.440\", \"../data/sb441-n16.bsp\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We create a simple rebound simulation and attach ASSIST."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Searching NASA Horizons for '2024 YR4'... \n",
+      "Found: (2024 YR4) \n",
+      "Initialized simulation at t=7304.5, initial position: [-0.5145660162497446, -2.8176650229161004, -1.26053378304063]\n"
+     ]
+    }
+   ],
+   "source": [
+    "T0 = 2458849.5\n",
+    "T0_sim = T0 - ephem.jd_ref\n",
+    "\n",
+    "def setup_sim():\n",
+    "    sim = rebound.Simulation()\n",
+    "    ax = assist.Extras(sim, ephem)\n",
+    "    sim.add(\"2024 YR4\", date=f\"JD{T0}\")\n",
+    "    sim.t = T0_sim\n",
+    "    return sim, ax\n",
+    "\n",
+    "sim, ax = setup_sim()\n",
+    "\n",
+    "print(f\"Initialized simulation at t={sim.t}, initial position: {sim.particles[0].xyz}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Our endpoint will be N days in the future, selected by SIM_LEN."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "We'll work with the times [7304.5 7305.5 7306.5]...[10301.5 10302.5 10303.5]\n",
+      "Searching NASA Horizons for '2024 YR4'... \n",
+      "Found: (2024 YR4) \n",
+      "Finished standard method simulation...sample:\n",
+      "[[-0.51456602 -2.81766502 -1.26053378]\n",
+      " [-0.50724626 -2.81389356 -1.2584097 ]\n",
+      " [-0.49992162 -2.81009494 -1.25627347]\n",
+      " [-0.49259214 -2.80626905 -1.25412505]\n",
+      " [-0.48525789 -2.80241579 -1.2519644 ]]...\n"
+     ]
+    }
+   ],
+   "source": [
+    "SIM_LEN = 3000\n",
+    "T_END = T0_sim + SIM_LEN\n",
+    "\n",
+    "times = np.linspace(T0_sim, T_END, SIM_LEN, endpoint=False)\n",
+    "print(f\"We'll work with the times {times[0:3]}...{times[-3:]}\")\n",
+    "positions_default = np.zeros((SIM_LEN, 3))\n",
+    "\n",
+    "sim1, ax1 = setup_sim()\n",
+    "for i, t in enumerate(times):\n",
+    "    ax1.integrate_or_interpolate(t)\n",
+    "    positions_default[i] = sim1.particles[0].xyz\n",
+    "\n",
+    "print(f\"Finished standard method simulation...sample:\\n{positions_default[0:5]}...\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Searching NASA Horizons for '2024 YR4'... \n",
+      "Found: (2024 YR4) \n"
+     ]
+    }
+   ],
+   "source": [
+    "sim2, ax2 = setup_sim()\n",
+    "sim2.save_to_file(\"sim2.bin\", step=1, delete_file=True)\n",
+    "ax2.integrate_or_interpolate(T_END+1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now test the file..."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Finished interpolating from SimulationArchive...sample:\n",
+      "[[-0.51456602 -2.81766502 -1.26053378]\n",
+      " [-0.50724626 -2.81389356 -1.2584097 ]\n",
+      " [-0.49992162 -2.81009494 -1.25627347]\n",
+      " [-0.49259214 -2.80626905 -1.25412505]\n",
+      " [-0.48525789 -2.80241579 -1.2519644 ]]...\n",
+      "\n",
+      "Average difference between the states: -3.77063658508342e-13\n"
+     ]
+    }
+   ],
+   "source": [
+    "sim_f = rebound.Simulation(\"sim2.bin\")\n",
+    "ax_f = assist.Extras(sim_f, ephem)\n",
+    "\n",
+    "positions_fromfile = np.zeros((SIM_LEN, 3))\n",
+    "\n",
+    "for i, t in enumerate(times):\n",
+    "    ax_f.integrate_or_interpolate(t)\n",
+    "    positions_fromfile[i] = sim_f.particles[0].xyz\n",
+    "\n",
+    "print(f\"Finished interpolating from SimulationArchive...sample:\\n{positions_fromfile[0:5]}...\\n\")\n",
+    "print(f\"Average difference between the states: {np.average(positions_default - positions_fromfile)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Average difference between the states: -8.970580559789846e-17\n"
+     ]
+    }
+   ],
+   "source": [
+    "from rebound.simulationarchive import Simulationarchive\n",
+    "sa=Simulationarchive(\"sim2.bin\")\n",
+    "\n",
+    "positions_fromfile_fixed = np.zeros((SIM_LEN-1,3))\n",
+    "for i, t in enumerate(times[1:]):\n",
+    "    tempsim = assist.tools.assist_create_interpolated_simulation(sa, t)\n",
+    "    positions_fromfile_fixed[i] = tempsim.particles[0].xyz\n",
+    "\n",
+    "print(f\"Average difference between the states: {np.average(positions_default[1:] - positions_fromfile_fixed)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note the difference from the C function is better, while the naive loading of the archive is wrong (this is for the extreme case of 2024 YR4, but nonetheless true).\n",
+    "\n",
+    "Also note that because we need the acceleration from a step prior, we cannot interpolate during the first timestep (this is hard-coded in the C function, though I'm not sure if it's because it's impossible or inconvenient/incompatible with out SimulationArchives are saved currently)."
+   ]
+  }
+ ],
+ "metadata": {
+  "anaconda-cloud": {},
+  "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.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}