{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "## Classical mechanics problems" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Prerequisites" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "def ode_euler_step(f, x, t, h):\n", " \"\"\"Perform a single step h using Euler's scheme.\n", "\n", " Args:\n", " f: the function that defines the ODE.\n", " x: the value of the dependent variable at the present step.\n", " t: the present value of the time variable.\n", " h: the time step\n", "\n", " Returns:\n", " xnew: the value of the dependent variable at the step t+h\n", " \"\"\"\n", " return x + h * f(x,t)\n", "\n", "def ode_rk4_step(f, x, t, h):\n", " \"\"\"Perform a single step h using 4th order Runge-Kutta method.\n", "\n", " Args:\n", " f: the function that defines the ODE.\n", " x: the value of the dependent variable at the present step.\n", " t: the present value of the time variable.\n", " h: the time step\n", "\n", " Returns:\n", " xnew: the value of the dependent variable at the step t+h\n", " \"\"\"\n", " k1 = h * f(x,t)\n", " k2 = h * f(x + k1/2., t + h /2.)\n", " k3 = h * f(x + k2/2., t + h /2.)\n", " k4 = h * f(x + k3, t + h)\n", " return x + (k1 + 2. * k2 + 2. * k3 + k4) / 6.\n", "\n", "# The default definition of the error (distance) between two state vectors \n", "# Default: the magnitude of the difference vector\n", "def distance_definition_default(x1, x2):\n", " diff = x1 - x2\n", " diffnorm = np.sqrt(np.dot(diff, diff))\n", " return diffnorm\n", "\n", "def ode_rk4_adaptive_multi(f, x0, t0, h0, tmax, delta = 1.e-6, distance_definition = distance_definition_default):\n", " \"\"\"Solve an ODE dx/dt = f(x,t) from t = t0 to t = t0 + h*steps \n", " using 4th order Runge-Kutta method with adaptive time step.\n", "\n", " Args:\n", " f: the function that defines the ODE.\n", " x0: the initial value of the dependent variable.\n", " t0: the initial value of the time variable.\n", " h0: the initial time step\n", " tmax: the maximum time\n", " delta: the desired accuracy per unit time\n", "\n", " Returns:\n", " t,x: the pair of arrays corresponding to the time and dependent variables\n", " \"\"\"\n", " \n", " ts = [t0]\n", " xs = [x0]\n", " \n", " h = h0\n", " t = t0\n", " i = 0\n", " \n", " while (t < tmax):\n", " if (t + h >= tmax):\n", " ts.append(tmax)\n", " h = tmax - t\n", " xs.append(ode_rk4_step(f, xs[i], ts[i], h))\n", " t = tmax\n", " break\n", " \n", " x1 = ode_rk4_step(f, xs[i], ts[i], h)\n", " x1 = ode_rk4_step(f, x1, ts[i] + h, h)\n", " x2 = ode_rk4_step(f, xs[i], ts[i], 2*h)\n", " \n", " diffnorm = distance_definition(x1, x2)\n", " if diffnorm == 0.: # To avoid the division by zero\n", " rho = 2.**4\n", " else:\n", " rho = 30. * h * delta / diffnorm\n", " if rho < 1.:\n", " h *= rho**(1/4.)\n", " else:\n", " if (t + 2.*h) < tmax:\n", " xs.append(x1)\n", " ts.append(t + 2*h)\n", " t += 2*h\n", " else:\n", " xs.append(ode_rk4_step(f, xs[i], ts[i], h))\n", " ts.append(t + h)\n", " t += h\n", " i += 1\n", " h = min(2.*h, h * rho**(1/4.))\n", " \n", " return ts,xs\n", "\n", "def ode_leapfrog_step(f, x, x2, t, h):\n", " \"\"\"Perform a single step h using the leapfrog method.\n", "\n", " Args:\n", " f: the function that defines the ODE.\n", " x: the value of x(t)\n", " x2: the value of x(t+h/2)\n", " t: the present value of the time variable.\n", " h: the time step\n", "\n", " Returns:\n", " xnew, xnew2: the value of the dependent variable at the steps t+h, t+3h/2\n", " \"\"\"\n", " \n", " xnew = x + h * f(x2,t+h/2.)\n", " xnew2 = x2 + h * f(xnew, t + h)\n", " return xnew, xnew2\n", "\n", "def ode_MMM_multi(f, x0, t0, H, nsteps):\n", " \"\"\"Multi-dimensional version of the modified midpoint method.\n", " \"\"\"\n", " \n", " h = H / nsteps\n", " t = np.zeros(nsteps + 1)\n", " x = np.zeros((len(t), len(x0)))\n", " x2 = np.zeros(len(x0))\n", " t = t0\n", " x = x0\n", " y = ode_euler_step(f, x0, t0, h/2.)\n", " for i in range(0, nsteps):\n", " yprev = y\n", " x, y = ode_leapfrog_step(f, x, y, t, h)\n", " t = t + h\n", " \n", " return 0.5 * (x + yprev + 0.5 * h * f(x,t))\n", "\n", "def bulirsch_stoer_step(f, x0, t0, H, delta = 1.e-6, distance_definition = distance_definition_default, maxsteps = 10):\n", " \"\"\"Use Bulirsch-Stoer method to integrate for t to t+H.\n", " \"\"\"\n", " n = 1\n", " R1 = np.empty([1,len(x0)],float)\n", " R1[0] = ode_MMM_multi(f, x0, t0, H, 1)\n", " error = 2. * H * delta\n", " while (error > H*delta or n < 2) and n < maxsteps:\n", " n += 1\n", " R2 = R1\n", " R1 = np.empty([n,len(x0)],float)\n", " R1[0] = ode_MMM_multi(f, x0, t0, H, n)\n", " for m in range(1,n):\n", " epsilon = (R1[m-1]-R2[m-1])/((n/(n-1))**(2*m)-1)\n", " R1[m] = R1[m-1] + epsilon\n", " error = distance_definition(R1[n-2],R1[n-1])\n", " \n", " if n == maxsteps:\n", " # Reached maximum number of substeps in Bulirsch-Stoer method\n", " # reducing the time step and applying the method recursively\n", " sol1 = bulirsch_stoer_step(f, x0, t0, H/2., delta, distance_definition, maxsteps)\n", " sol2 = bulirsch_stoer_step(f, sol1[-1][1], t0 + H/2., H/2., delta, distance_definition, maxsteps)\n", " return sol1 + sol2\n", " \n", " return [[t0+H, R1[n - 1]]]\n", "\n", "def bulirsch_stoer(f, x0, t0, nsteps, tmax, delta = 1.e-6, distance_definition = distance_definition_default, maxsubsteps = 10):\n", " \"\"\"Use Bulirsch-Stoer method to integrate for t to tmax using nsteps Bulirsch-Stoer steps\n", " \"\"\"\n", " H = (tmax - t0) / nsteps\n", " t = np.zeros(nsteps + 1)\n", " x = np.zeros((len(t), len(x0)))\n", " t = [t0]\n", " x = [x0]\n", " for i in range(0, nsteps):\n", " bst = bulirsch_stoer_step(f, x[-1], t[-1], H, delta, distance_definition, maxsubsteps)\n", " [t.append(el[0]) for el in bst]\n", " [x.append(el[1]) for el in bst]\n", " return t,x" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Three-body problem\n", "\n", "Exercise 8.16 from M. Newman *Computational Physics*\n", "\n", "We have three stars initially at rest and interacting through gravitational force.\n", "The dimensionless parameters are the following:\n", "\n", "| Name | Mass | x | y|\n", "| ----------- | ----------- |----------- | ----------- |\n", "| Star 1 | 150. | 3 | 1 |\n", "| Star 2 | 200. | -1 | -2|\n", "| Star 3 | 250. | -1 | 1\n", "\n", "The motion is in $z = 0$ plane.\n", "The gravitation constant is taken to be $G = 1$\n", "\n", "The equations of motion are\n", "\\begin{align*}\n", "\\frac{d^2 \\mathbf{r}_1}{dt} &= G m_2 \\frac{\\mathbf{r}_2-\\mathbf{r}_1}{|\\mathbf{r}_2-\\mathbf{r}_1|^3} + G m_3 \\frac{\\mathbf{r}_3-\\mathbf{r}_1}{|\\mathbf{r}_3-\\mathbf{r}_1|^3},\\\\\n", "\\frac{d^2 \\mathbf{r}_2}{dt} &= G m_1 \\frac{\\mathbf{r}_1-\\mathbf{r}_2}{|\\mathbf{r}_1-\\mathbf{r}_2|^3} + G m_3 \\frac{\\mathbf{r}_3-\\mathbf{r}_2}{|\\mathbf{r}_3-\\mathbf{r}_2|^3},\\\\\n", "\\frac{d^2 \\mathbf{r}_3}{dt} &= G m_1 \\frac{\\mathbf{r}_1-\\mathbf{r}_3}{|\\mathbf{r}_1-\\mathbf{r}_3|^3} + G m_2 \\frac{\\mathbf{r}_2-\\mathbf{r}_3}{|\\mathbf{r}_2-\\mathbf{r}_3|^3}.\n", "\\end{align*}" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "f_evaluations = 0\n", "\n", "def fthreebody(xin, t):\n", " global f_evaluations\n", " f_evaluations += 1\n", " \n", " x1 = xin[0]\n", " y1 = xin[1]\n", " x2 = xin[2]\n", " y2 = xin[3]\n", " x3 = xin[4]\n", " y3 = xin[5]\n", " \n", " r12 = np.sqrt((x1-x2)**2 + (y1-y2)**2)\n", " r13 = np.sqrt((x1-x3)**2 + (y1-y3)**2)\n", " r23 = np.sqrt((x2-x3)**2 + (y2-y3)**2)\n", " \n", " \n", " return np.array([xin[6],xin[7],xin[8],xin[9],xin[10],xin[11],\n", " G * m2 * (x2 - x1) / r12**3 + G * m3 * (x3 - x1) / r13**3,\n", " G * m2 * (y2 - y1) / r12**3 + G * m3 * (y3 - y1) / r13**3,\n", " G * m1 * (x1 - x2) / r12**3 + G * m3 * (x3 - x2) / r23**3,\n", " G * m1 * (y1 - y2) / r12**3 + G * m3 * (y3 - y2) / r23**3,\n", " G * m1 * (x1 - x3) / r13**3 + G * m2 * (x2 - x3) / r23**3,\n", " G * m1 * (y1 - y3) / r13**3 + G * m2 * (y2 - y3) / r23**3\n", " ]\n", " )\n", "\n", "def error_definition_threebody(x1, x2):\n", " val = 0.\n", " for i in range(0,6):\n", " val += (x1[i] - x2[i])**2\n", " return np.sqrt(val)\n", "\n", "def threebody_kinetic_energy(xin):\n", " val = 0.\n", " ms = [m1,m1,m2,m2,m3,m3]\n", " for i in range(0,6):\n", " val += ms[i] * xin[6 + i]**2 / 2.\n", " return val\n", "\n", "def threebody_potential_energy(xin):\n", " x1 = xin[0]\n", " y1 = xin[1]\n", " x2 = xin[2]\n", " y2 = xin[3]\n", " x3 = xin[4]\n", " y3 = xin[5]\n", " \n", " r12 = np.sqrt((x1-x2)**2 + (y1-y2)**2)\n", " r13 = np.sqrt((x1-x3)**2 + (y1-y3)**2)\n", " r23 = np.sqrt((x2-x3)**2 + (y2-y3)**2)\n", " return -G * m1 * m2 / r12 - G * m1 * m3 / r13 - G * m2 * m3 / r23" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
" ], "text/plain": [ "