{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# \"The Crank-Nicolson method combined with Runge-Kutta implemented from scratch in Python\"\n",
    "> \"In this article we implement the well-known finite difference method Crank-Nicolson in combination with a Runge-Kutta solver in Python.\"\n",
    "- toc: true\n",
    "- branch: master\n",
    "- badges: true\n",
    "- comments: true\n",
    "- categories: [python, numpy, numerical analysis, partial differential equations]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combining Crank-Nicolson and Runge-Kutta to Solve a Reaction-Diffusion System"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have already [derived](http://georg.io/2013/12/Crank_Nicolson/) the Crank-Nicolson method\n",
    "to integrate the following reaction-diffusion system numerically:\n",
    "\n",
    "$$\\frac{\\partial u}{\\partial t} = D \\frac{\\partial^2 u}{\\partial x^2} + f(u),$$\n",
    "\n",
    "$$\\frac{\\partial u}{\\partial x}\\Bigg|_{x = 0, L} = 0.$$\n",
    "\n",
    "Please refer to the [earlier blog post](http://georg.io/2013/12/Crank_Nicolson/) for details.\n",
    "\n",
    "In our previous derivation, we constructed the following stencil that we would go on to\n",
    "rearrange into a system of linear equations that we needed to solve every time step:\n",
    "\n",
    "$$\\frac{U_j^{n+1} - U_j^n}{\\Delta t} = \\frac{D}{2 \\Delta x^2} \\left( U_{j+1}^n - 2 U_j^n + U_{j-1}^n + U_{j+1}^{n+1} - 2 U_j^{n+1} + U_{j-1}^{n+1}\\right) + f(U_j^n),$$\n",
    "\n",
    "where $j$ and $n$ are space and time grid points respectively.\n",
    "\n",
    "Rearranging the above set of equations, we effectively integrate the reaction part with the\n",
    "[explicit Euler method](https://en.wikipedia.org/wiki/Euler_method) like so:\n",
    "\n",
    "$$U_j^{n+1} = U_j^n + \\Delta t f(U_j^n).$$\n",
    "\n",
    "For functions $f$ that change rapidly for small changes in their input\n",
    "([stiff equations](https://en.wikipedia.org/wiki/Stiff_equation)), using\n",
    "the explicit Euler method may pose stability problems unless we choose a sufficiently\n",
    "small $\\Delta t$.\n",
    "\n",
    "Therefore, I have been wondering if it would be possible to use a more sophisticated\n",
    "and stable numerical scheme to integrate the reaction part in the context of our\n",
    "Crank-Nicolson scheme.\n",
    "\n",
    "For instance, to integrate the reaction part with the\n",
    "[classical Runge-Kutta method](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods#The_Runge.E2.80.93Kutta_method),\n",
    "we would write out the following set of equations instead of the aforementioned one:\n",
    "\n",
    "$$\\frac{U_j^{n+1} - U_j^n}{\\Delta t} = \\frac{D}{2 \\Delta x^2} \\left( U_{j+1}^n - 2 U_j^n + U_{j-1}^n + U_{j+1}^{n+1} - 2 U_j^{n+1} + U_{j-1}^{n+1}\\right) + \\frac{1}{6} \\left(k_1 + 2 k_2 + 2 k_3 + k_4 \\right),$$\n",
    "\n",
    "where\n",
    "\n",
    "$$k_1 = f(U_j^n),$$\n",
    "\n",
    "$$k_2 = f\\left( U_j^n + \\frac{\\Delta t}{2} k_1 \\right),$$\n",
    "\n",
    "$$k_3 = f\\left( U_j^n + \\frac{\\Delta t}{2} k_2 \\right),$$\n",
    "\n",
    "$$k_4 = f\\left( U_j^n + \\Delta t k_3 \\right).$$\n",
    "\n",
    "Whether or not doing this makes sense theoretically I am not certain. But going ahead and implementing this\n",
    "to the numerical example we discussed [earlier](http://georg.io/2013/12/Crank_Nicolson) seems to suggest\n",
    "that this does work.\n",
    "\n",
    "In the following Python code that is mostly a copy of [our previous code](http://georg.io/2013/12/Crank_Nicolson/)\n",
    "we compare the time behaviour and accuracy (measured by mass conservation as our reaction diffusion system\n",
    "preserves mass) of the explicit Euler and Runge-Kutta 4 reaction integration.\n",
    "\n",
    "We realize that the differences between the obtained numerical results are negligible and we shall\n",
    "compare both approaches with a stiffer reaction term another time.\n",
    "\n",
    "We shall also take a look at more sophisticated measures of numerical stability another time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy\n",
    "from matplotlib import pyplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numpy.set_printoptions(precision=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 1.\n",
    "J = 200\n",
    "dx = float(L)/float(J-1)\n",
    "x_grid = numpy.array([j*dx for j in range(J)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 500\n",
    "N = 1000\n",
    "dt = float(T)/float(N-1)\n",
    "t_grid = numpy.array([n*dt for n in range(N)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D_v = float(10.)/float(100.)\n",
    "D_u = 0.01 * D_v\n",
    "\n",
    "k0 = 0.067\n",
    "f = lambda u, v: dt*(v*(k0 + float(u*u)/float(1. + u*u)) - u)\n",
    "g = lambda u, v: -f(u,v)\n",
    " \n",
    "sigma_u = float(D_u*dt)/float((2.*dx*dx))\n",
    "sigma_v = float(D_v*dt)/float((2.*dx*dx))\n",
    "\n",
    "total_protein = 2.26"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_high = 10\n",
    "U =  numpy.array([0.1 for i in range(no_high,J)] + [2. for i in range(0,no_high)])\n",
    "V = numpy.array([float(total_protein-dx*sum(U))/float(J*dx) for i in range(0,J)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us take a look at the inhomogeneous initial condition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyplot.ylim((0., 2.1))\n",
    "pyplot.xlabel('x')\n",
    "pyplot.ylabel('concentration')\n",
    "pyplot.plot(x_grid, U)\n",
    "pyplot.plot(x_grid, V)\n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the matrices of our system of linear equations whose derivation \n",
    "[we described earlier](http://georg.io/2013/12/Crank_Nicolson/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_u = numpy.diagflat([-sigma_u for i in range(J-1)], -1) +\\\n",
    "      numpy.diagflat([1.+sigma_u]+[1.+2.*sigma_u for i in range(J-2)]+[1.+sigma_u]) +\\\n",
    "      numpy.diagflat([-sigma_u for i in range(J-1)], 1)\n",
    "        \n",
    "B_u = numpy.diagflat([sigma_u for i in range(J-1)], -1) +\\\n",
    "      numpy.diagflat([1.-sigma_u]+[1.-2.*sigma_u for i in range(J-2)]+[1.-sigma_u]) +\\\n",
    "      numpy.diagflat([sigma_u for i in range(J-1)], 1)\n",
    "        \n",
    "A_v = numpy.diagflat([-sigma_v for i in range(J-1)], -1) +\\\n",
    "      numpy.diagflat([1.+sigma_v]+[1.+2.*sigma_v for i in range(J-2)]+[1.+sigma_v]) +\\\n",
    "      numpy.diagflat([-sigma_v for i in range(J-1)], 1)\n",
    "        \n",
    "B_v = numpy.diagflat([sigma_v for i in range(J-1)], -1) +\\\n",
    "      numpy.diagflat([1.-sigma_v]+[1.-2.*sigma_v for i in range(J-2)]+[1.-sigma_v]) +\\\n",
    "      numpy.diagflat([sigma_v for i in range(J-1)], 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function `f_vec_ee` returns the explicit Euler time step vector while `f_vec_rk` returns the vector obtained\n",
    "applying the Runge-Kutta 4 method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_vec_ee(U,V):\n",
    "    return numpy.multiply(dt, numpy.subtract(numpy.multiply(V, \n",
    "           numpy.add(k0, numpy.divide(numpy.multiply(U,U), numpy.add(1., numpy.multiply(U,U))))), U))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_vec_rk(U, V):\n",
    "    f_vec = lambda U, V: numpy.subtract(numpy.multiply(V, \n",
    "                         numpy.add(k0, numpy.divide(numpy.multiply(U,U), numpy.add(1., numpy.multiply(U,U))))), U)\n",
    "    k1 = f_vec(U, V)\n",
    "    k2 = f_vec(U + numpy.multiply(dt/2., k1), V - numpy.multiply(dt/2., k1))\n",
    "    k3 = f_vec(U + numpy.multiply(dt/2., k2), V - numpy.multiply(dt/2., k2))\n",
    "    k4 = f_vec(U + numpy.multiply(dt, k3), V - numpy.multiply(dt, k3))\n",
    "    \n",
    "    return numpy.multiply(dt/6., k1 + numpy.multiply(2., k2) + numpy.multiply(2., k3) + k4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U_record_ee = numpy.empty(shape=(N,J))\n",
    "V_record_ee = numpy.empty(shape=(N,J))\n",
    "\n",
    "U_record_rk = numpy.empty(shape=(N,J))\n",
    "V_record_rk = numpy.empty(shape=(N,J))\n",
    "\n",
    "U_record_ee[0][:] = U[:]\n",
    "V_record_ee[0][:] = V[:]\n",
    "\n",
    "U_record_rk[0][:] = U[:]\n",
    "V_record_rk[0][:] = V[:]\n",
    "\n",
    "for ti in range(1,N):\n",
    "    U_record_ee[ti][:] = numpy.linalg.solve(A_u, B_u.dot(U_record_ee[ti-1][:]) +\n",
    "                                            f_vec_ee(U_record_ee[ti-1][:],V_record_ee[ti-1][:]))\n",
    "    V_record_ee[ti][:] = numpy.linalg.solve(A_v, B_v.dot(V_record_ee[ti-1][:]) - \n",
    "                                            f_vec_ee(U_record_ee[ti-1][:],V_record_ee[ti-1][:]))\n",
    "    \n",
    "    U_record_rk[ti][:] = numpy.linalg.solve(A_u, B_u.dot(U_record_rk[ti-1][:]) +\n",
    "                                            f_vec_rk(U_record_rk[ti-1][:],V_record_rk[ti-1][:]))\n",
    "    V_record_rk[ti][:] = numpy.linalg.solve(A_v, B_v.dot(V_record_rk[ti-1][:]) - \n",
    "                                            f_vec_rk(U_record_rk[ti-1][:],V_record_rk[ti-1][:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial protein mass in our system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print 'Explicit Euler', numpy.sum(numpy.multiply(dx, U_record_ee[0]) + numpy.multiply(dx, V_record_ee[0]))\n",
    "print 'Runge-Kutta 4 ', numpy.sum(numpy.multiply(dx, U_record_ee[0]) + numpy.multiply(dx, V_record_ee[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since our reaction-diffusion system preserves mass, we should retain the same protein mass at steady-state\n",
    "for both numerical approaches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print 'Explicit Euler %.14f' % numpy.sum(numpy.multiply(dx, U_record_ee[-1]) + numpy.multiply(dx, V_record_ee[-1]))\n",
    "print 'Runge-Kutta 4  %.14f' % numpy.sum(numpy.multiply(dx, U_record_rk[-1]) + numpy.multiply(dx, V_record_rk[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We realize that the difference between the two numerical methods is neglibible and we shall\n",
    "compare both approaches for a stiffer system another time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A plot of the steady-state concentration profiles confirms that we cannot observe a significant differences\n",
    "between the results generated by both methods (varying `J` and `N` paints the same pictures)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyplot.ylim((0., 2.1))\n",
    "pyplot.xlabel('x')\n",
    "pyplot.ylabel('concentration')\n",
    "pyplot.plot(x_grid, U_record_ee[-1])\n",
    "pyplot.plot(x_grid, V_record_ee[-1])\n",
    "pyplot.plot(x_grid, U_record_rk[-1])\n",
    "pyplot.plot(x_grid, V_record_rk[-1])\n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kymograph of `U` integrated with the explicit Euler method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = pyplot.subplots()\n",
    "pyplot.xlabel('x')\n",
    "pyplot.ylabel('t')\n",
    "pyplot.ylim((0., T))\n",
    "heatmap = ax.pcolormesh(x_grid, t_grid, U_record_ee, vmin=0., vmax=1.2)\n",
    "colorbar = pyplot.colorbar(heatmap)\n",
    "colorbar.set_label('concentration U')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kymograph of `U` integrated with the Runge-Kutta 4 method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = pyplot.subplots()\n",
    "pyplot.xlabel('x')\n",
    "pyplot.ylabel('t')\n",
    "pyplot.ylim((0., T))\n",
    "heatmap = ax.pcolormesh(x_grid, t_grid, U_record_rk, vmin=0., vmax=1.2)\n",
    "colorbar = pyplot.colorbar(heatmap)\n",
    "colorbar.set_label('concentration U')"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}