Relaxation_RK_reversibility.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ff384b7-1d40-4846-a01f-82bda946d9b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from nodepy import rk, stability_function\n",
    "from scipy.optimize import root, fsolve, newton, brentq, bisect\n",
    "\n",
    "rk4 = rk.loadRKM('RK44').__num__()\n",
    "rk4x2 = rk4*rk4\n",
    "ssp2 = rk.loadRKM('SSP22').__num__()\n",
    "ssp3 = rk.loadRKM('SSP33').__num__()\n",
    "ssp104 = rk.loadRKM('SSP104').__num__()\n",
    "merson4 = rk.loadRKM('Merson43').__num__()\n",
    "bs5 = rk.loadRKM('BS5').__num__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e30307e8-d51e-4ba1-b541-9708270cf8bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def RRK_quadratic_one_step(rkm, dt, f, w0=[1.,0], t_final=10., relaxation=True, \n",
    "        rescale_step=True, debug=False, gammatol=0.1):\n",
    "    \"\"\"\n",
    "    Relaxation Runge-Kutta method implementation for quadratic energy functional.\n",
    "    This function takes a single step.\n",
    "    \n",
    "    Options:\n",
    "    \n",
    "        rkm: Base Runge-Kutta method, in Nodepy format\n",
    "        dt: time step size\n",
    "        f: RHS of ODE system\n",
    "        w0: Initial data\n",
    "        t_final: final solution time\n",
    "        relaxation: if True, use relaxation method.  Otherwise, use vanilla RK method.\n",
    "        rescale_step: if True, new time step is t_n + \\gamma dt\n",
    "        debug: output some additional diagnostics\n",
    "        gammatol: Fail if abs(1-gamma) exceeds this value\n",
    "        \n",
    "    \"\"\"\n",
    "    w = np.array(w0)\n",
    "    t = 0\n",
    "    s = len(rkm)\n",
    "    b = rkm.b\n",
    "    y = np.zeros((s,len(w0)))\n",
    "    \n",
    "    for i in range(s):\n",
    "        y[i,:] = w.copy()\n",
    "        for j in range(i):\n",
    "            y[i,:] += rkm.A[i,j]*dt*f(y[j,:])\n",
    "            \n",
    "    F = np.array([f(y[i,:]) for i in range(s)])\n",
    "    \n",
    "    if relaxation:\n",
    "        numer = 2*sum(b[i]*rkm.A[i,j]*np.dot(F[i],F[j]) \\\n",
    "                            for i in range(s) for j in range(s))\n",
    "        denom = sum(b[i]*b[j]*np.dot(F[i],F[j]) for i in range(s) for j in range(s))\n",
    "        if denom != 0:\n",
    "            gam = numer/denom\n",
    "        else:\n",
    "            gam = 1.\n",
    "    else:  # Use standard RK method\n",
    "        gam = 1.\n",
    "    \n",
    "    w = w + gam*dt*sum([b[j]*F[j] for j in range(s)])\n",
    "        \n",
    "    return w, gam\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7198aef-10f9-4479-ba1d-2af74ca91e42",
   "metadata": {},
   "source": [
    "# Testing with the harmonic oscillator and Ranocha's oscillator "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "12e4f2b6-5df2-4176-9044-6ee6cf7c94d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_HO(u):  # harmonic oscillator\n",
    "    return np.array([-u[1],u[0]])\n",
    "    \n",
    "def f_ranocha(u):\n",
    "    E = u[0]**2 + u[1]**2\n",
    "    return np.array([-u[1],u[0]])/E"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb8997b2-644b-4cac-8662-7bb29c8e51f3",
   "metadata": {},
   "source": [
    "If you choose the perfect step size, even standard RK methods can be reversible for\n",
    "the harmonic oscillator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c3ff7165-b6f0-4433-973a-049b17866468",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-4.44089210e-16 -2.22044605e-16]\n"
     ]
    }
   ],
   "source": [
    "method = ssp3; dt = np.sqrt(3)\n",
    "ff = f_HO; u0 = np.array([1.,0.])\n",
    "\n",
    "u1, gam = RRK_quadratic_one_step(method, dt, ff, w0=u0, relaxation=False)\n",
    "u1[1] = -u1[1]   # p -> -p\n",
    "u2, gam = RRK_quadratic_one_step(method, dt, ff, w0=u1, relaxation=False)\n",
    "print(u2-u0)  # this will vanish if the method is reversible"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6f4ee39-681c-4767-b382-3fcbea296f9b",
   "metadata": {},
   "source": [
    "But not for more complicated problems:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "58156be1-a847-4ee2-acb1-e6fffee4be40",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.28242453 -0.72696541]\n"
     ]
    }
   ],
   "source": [
    "method = ssp3; dt = np.sqrt(3)\n",
    "ff = f_ranocha; u0 = np.array([1.,0.])\n",
    "\n",
    "u1, gam = RRK_quadratic_one_step(method, dt, ff, w0=u0, relaxation=False)\n",
    "u1[1] = -u1[1]   # p -> -p\n",
    "u2, gam = RRK_quadratic_one_step(method, dt, ff, w0=u1, relaxation=False)\n",
    "print(u2-u0)  # this will vanish if the method is reversible"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e61ab31-7337-4fcb-993c-75d623b2f204",
   "metadata": {},
   "source": [
    "With relaxation, every RK method is reversible under any step size, for both the linear and \n",
    "the nonlinear oscillator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b6fb7559-7df1-4f4e-8d7a-e8012b1228d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.00000000e+00 -1.11022302e-16]\n"
     ]
    }
   ],
   "source": [
    "method = ssp3; dt = 0.5\n",
    "ff = f_HO; u0 = np.array([1.,0.])\n",
    "\n",
    "u1, gam = RRK_quadratic_one_step(method, dt, ff, w0=u0, relaxation=True)\n",
    "u1[1] = -u1[1]   # p -> -p\n",
    "u2, gam = RRK_quadratic_one_step(method, dt, ff, w0=u1, relaxation=True)\n",
    "print(u2-u0)  # this will vanish if the method is reversible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fe5a14b5-493b-487f-80b7-0a4a4ee6b226",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.00000000e+00 -5.55111512e-17]\n"
     ]
    }
   ],
   "source": [
    "method = ssp3; dt = 0.5\n",
    "ff = f_ranocha; u0 = np.array([1.,0.])\n",
    "\n",
    "u1, gam = RRK_quadratic_one_step(method, dt, ff, w0=u0, relaxation=True)\n",
    "u1[1] = -u1[1]   # p -> -p\n",
    "u2, gam = RRK_quadratic_one_step(method, dt, ff, w0=u1, relaxation=True)\n",
    "print(u2-u0)  # this will vanish if the method is reversible"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd9e6893-ac3c-49dd-b8cb-17721322fb84",
   "metadata": {},
   "source": [
    "# Kepler problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "042a4dd4-9b2f-4c70-82a0-3ba57c7b796f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# u = [q_1 q_2 p_1 p_2]\n",
    "\n",
    "def f_kepler(u):\n",
    "    q3 = (u[0]**2+u[1]**2)**(3/2.)\n",
    "    return np.array([u[2], u[3], -u[0]/q3, -u[1]/q3])\n",
    "\n",
    "def hamiltonian(w):\n",
    "    abs2_q = w[0]*w[0] + w[1]*w[1]\n",
    "    abs2_p = w[2]*w[2] + w[3]*w[3]\n",
    "    return 0.5 * abs2_p - 1.0 / np.sqrt(abs2_q)\n",
    "\n",
    "def d_hamiltonian(w):\n",
    "    q1 = w[0]\n",
    "    q2 = w[1]\n",
    "    p1 = w[2]\n",
    "    p2 = w[3]\n",
    "    abs_q = np.sqrt(q1*q1 + q2*q2)\n",
    "    \n",
    "    dq1 = q1 / (abs_q*abs_q*abs_q)\n",
    "    dq2 = q2 / (abs_q*abs_q*abs_q)\n",
    "    dp1 = p1\n",
    "    dp2 = p2\n",
    "    return np.array([dq1, dq2, dp1, dp2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "43e8b966-a233-4f3f-979b-a29089820f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def RRK_one_step(rkm, dt, f, eta, deta, w0=[1.,0], t_final=10., relaxation=True, \n",
    "                 rescale_step=True, debug=False, gammatol=0.1, solver=\"brentq\", tol=1.e-14,\n",
    "                 maxiter=10000):\n",
    "    \"\"\"\n",
    "    Relaxation Runge-Kutta method implementation for general energy functional.\n",
    "    This function takes a single step.\n",
    "    \n",
    "    Options:\n",
    "    \n",
    "        rkm: Base Runge-Kutta method, in Nodepy format\n",
    "        dt: time step size\n",
    "        f: RHS of ODE system\n",
    "        w0: Initial data\n",
    "        t_final: final solution time\n",
    "        relaxation: if True, use relaxation method.  Otherwise, use vanilla RK method.\n",
    "        rescale_step: if True, new time step is t_n + \\gamma dt\n",
    "        debug: output some additional diagnostics\n",
    "        gammatol: Fail if abs(1-gamma) exceeds this value\n",
    "        \n",
    "    \"\"\"\n",
    "    w = np.array(w0)\n",
    "    t = 0\n",
    "    s = len(rkm)\n",
    "    b = rkm.b\n",
    "    y = np.zeros((s,len(w0)))\n",
    "    old_gamma = 1.0\n",
    "    \n",
    "    for i in range(s):\n",
    "        y[i,:] = w.copy()\n",
    "        for j in range(i):\n",
    "            y[i,:] += rkm.A[i,j]*dt*f(y[j,:])\n",
    "            \n",
    "    F = np.array([f(y[i,:]) for i in range(s)])\n",
    "    \n",
    "    if relaxation:\n",
    "        direction = dt * sum([b[i]*F[i,:] for i in range(s)])\n",
    "\n",
    "        r = lambda gamma: eta(w+gamma*direction) - eta(w)\n",
    "        rjac = lambda gamma: np.array([np.dot(deta(w+gamma*direction), direction)])\n",
    "\n",
    "        if solver == \"newton\":\n",
    "            gam = newton(r, old_gamma, fprime=rjac, tol=tol, maxiter=maxiter)\n",
    "            success = True\n",
    "            msg = \"Newton method did not converge\"\n",
    "        elif solver == \"brentq\" or solver == \"bisect\":\n",
    "            left = 0.9 * old_gamma\n",
    "            right = 1.1 * old_gamma\n",
    "            left_right_iter = 0\n",
    "            while r(left) * r(right) > 0:    \n",
    "                left *= 0.9\n",
    "                right *= 1.1\n",
    "                left_right_iter += 1\n",
    "                if left_right_iter > 100:\n",
    "                    raise SolveForGammaException(\n",
    "                        \"No suitable bounds found after %d iterations.\\nLeft = %e; r(left) = %e\\nRight = %e; r(right) = %e\\n\"%(\n",
    "                            left_right_iter, left, r(left), right, r(right)),w)\n",
    "\n",
    "            if solver == \"brentq\":\n",
    "                gam = brentq(r, left, right, xtol=tol, rtol=tol, maxiter=maxiter)\n",
    "            else:\n",
    "                gam = bisect(r, left, right, xtol=tol, rtol=tol, maxiter=maxiter)\n",
    "            success = True\n",
    "            msg = \"%s method did not converge\"%solver\n",
    "        else:\n",
    "            sol = root(r, old_gamma, jac=use_jac, method=solver, tol=tol,\n",
    "                       options={'xtol': tol, 'maxiter': maxiter})\n",
    "            gam = sol.x; success = sol.success; msg = sol.message\n",
    "\n",
    "    else:  # Use standard RK method\n",
    "        gam = 1.\n",
    "    \n",
    "    w = w + gam*dt*sum([b[j]*F[j] for j in range(s)])\n",
    "        \n",
    "    return w, gam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "29eaa0b6-b505-4acd-a871-e4d422b27e87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.5        0.         0.         0.57735027]\n",
      "[ 0.4949874   0.02877074 -0.20100885  0.57151346] 1.000002144903433\n",
      "-1.8333333333333333 -1.8333333333333335\n"
     ]
    }
   ],
   "source": [
    "method = rk4\n",
    "dt = 0.05\n",
    "solver = \"brentq\"\n",
    "e=0.5; u0 = np.array([1-e,0,0,np.sqrt((1-e)/(1+e))]); print(u0);\n",
    "\n",
    "u1, gam = RRK_one_step(method, dt, f_kepler, hamiltonian, d_hamiltonian, w0=u0, relaxation=True, t_final=100, solver=solver)\n",
    "print(u1,gam)\n",
    "print(hamiltonian(u0),hamiltonian(u1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "097cb4f2-be06-4848-9dcd-836616b4aca9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 4.99999949e-01 -7.47473992e-07 -5.34782400e-06 -5.77350626e-01] 1.0000279077174952\n"
     ]
    }
   ],
   "source": [
    "u1[2] = -u1[2]; u1[3]=-u1[3];\n",
    "u2, gam = RRK_one_step(method, dt, f_kepler, hamiltonian, d_hamiltonian, w0=u1, relaxation=True, t_final=100, solver=solver)\n",
    "print(u2,gam)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0733275b-dd86-4dcb-81ee-0076528dbd44",
   "metadata": {},
   "source": [
    "The method seems to be reversible, up to the accuracy of the nonlinear algebraic solve."
   ]
  }
 ],
 "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}