template_RK4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Runge kutta:\n",
    "# https://www.pp.rhul.ac.uk/~cowan/ph2150/kepler_xy.pdf\n",
    "\n",
    "class Simulator:\n",
    "    def __init__(self, initial_conditions, parameters, N_steps, dt) -> None:\n",
    "        self.params = parameters\n",
    "        self.N_steps = N_steps\n",
    "        self.dt = dt\n",
    "\n",
    "        # Array with time and array with coordinates and derivatives\n",
    "        self.t=np.arange(0,N_steps*self.dt,self.dt)\n",
    "        self.u = np.zeros((N_steps, len(initial_conditions)))\n",
    "\n",
    "        self.u[0] = np.array(initial_conditions)\n",
    "\n",
    "        self.runge_kutta()\n",
    "    \n",
    "    def ODE(self, t, u): # U = q1, q2, ... qn , dot_q1, dot_q2, ... , dot_qn\n",
    "        dot_u = np.zeros(u.size)\n",
    "        params = self.params\n",
    "        # ========= WRITE ODE ===============\n",
    "        # dot_u[0] = u[n]\n",
    "        # dot_u[1] = u[n+1]\n",
    "        # (...)\n",
    "        # dot_u[n] = f(u, params)\n",
    "        # (...)\n",
    "        # ====================================\n",
    "        return dot_u\n",
    "    \n",
    "    # Run runge kutta algorithm and save results in self.u\n",
    "    def runge_kutta(self):\n",
    "        for i in range(self.N_steps-1):\n",
    "            k1 = self.dt*self.ODE(self.dt*i, self.u[i])\n",
    "            k2 = self.dt*self.ODE(self.dt*(i+0.5), self.u[i] + 0.5*k1)\n",
    "            k3 = self.dt*self.ODE(self.dt*(i+0.5), self.u[i] + 0.5*k2)\n",
    "            k4 = self.dt*self.ODE(self.dt, self.u[i] + k3)\n",
    "            self.u[i+1] = self.u[i] + k1/6 + k2/3 + k3/3 + k4/6\n",
    "\n",
    "    # Plot coordinates over time\n",
    "    def plot_data_vs_time(self):\n",
    "        data = np.transpose(self.u)\n",
    "        plt.plot(self.t, data[0])\n",
    "        plt.plot(self.t, data[1])\n",
    "        # (...) you can plot all n coordinates if you want\n",
    "        plt.show()\n",
    "        return"
   ]
  }
 ],
 "metadata": {
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}