{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fe01d49f9c6740ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:41.599150Z",
     "start_time": "2025-03-25T10:54:40.565727Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.integrate import odeint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a91fd48ddb1c4fc",
   "metadata": {},
   "source": [
    "## Exercise 2.18 : Numerical precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:41.610056Z",
     "start_time": "2025-03-25T10:54:41.606741Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def f(t, x):\n",
    "    return 2 * np.power(np.abs(x), 1 / 2)\n",
    "\n",
    "\n",
    "def f_exact(x, lam):\n",
    "    return np.power(x - lam, 2) * (x >= lam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f447a3bdff49e850",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:42.273536Z",
     "start_time": "2025-03-25T10:54:42.268267Z"
    }
   },
   "outputs": [],
   "source": [
    "def euler_explicit(N, T, y0, f, type=int, args=()):\n",
    "    h = T / N\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    y = np.zeros(N + 1, dtype=type)\n",
    "    y[0] = y0\n",
    "    for n in range(N):\n",
    "        y[n + 1] = y[n] + h * f(t[n], y[n], *args)\n",
    "    return t, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "957b7a23b6076f77",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:42.767856Z",
     "start_time": "2025-03-25T10:54:42.638591Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17648 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "T = 10\n",
    "N = 1000\n",
    "u0 = 0\n",
    "\n",
    "t, sol_appr = euler_explicit(N, T, u0, f)\n",
    "x = np.linspace(u0, T, 100)\n",
    "\n",
    "plt.scatter(t, sol_appr, label=\"Approximation with EE\")\n",
    "plt.plot(x, f_exact(x, T), label=\"Exact solution\", color=\"red\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "57527c2e5eb6d17d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:43.351276Z",
     "start_time": "2025-03-25T10:54:43.345778Z"
    }
   },
   "outputs": [],
   "source": [
    "def euler_implicit(N, T, y0):\n",
    "    h = T / N\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    y = np.zeros(N + 1)\n",
    "    y[0] = y0\n",
    "    for n in range(N):\n",
    "        y[n + 1] = np.power(h + np.sqrt(h**2 + y[n]), 2)\n",
    "    return t, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2ff29a8872a51c89",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:44.080895Z",
     "start_time": "2025-03-25T10:54:43.982842Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35000 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "T = 10\n",
    "N = 1000\n",
    "u0 = 0\n",
    "\n",
    "t, sol_appr = euler_implicit(N, T, u0)\n",
    "x = np.linspace(0, T, 100)\n",
    "\n",
    "plt.scatter(t, sol_appr, label=\"Approximation with EI\")\n",
    "plt.plot(x, f_exact(x, T), label=\"Exact solution\", color=\"red\")\n",
    "plt.plot(x, x**2, label=\"Square function\", color=\"green\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecea2cd935dea93a",
   "metadata": {},
   "source": [
    "## Exercise 2.19 : SIR model, schema order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "61588e94b56ea779",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:45.274127Z",
     "start_time": "2025-03-25T10:54:45.270681Z"
    }
   },
   "outputs": [],
   "source": [
    "def F(y, t, a, r):\n",
    "    S, I, R = y\n",
    "    dS = -r * S * I\n",
    "    dI = r * S * I - a * I\n",
    "    dR = a * I\n",
    "    return np.array([dS, dI, dR])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a32f0facb958a1a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:45.743525Z",
     "start_time": "2025-03-25T10:54:45.739660Z"
    }
   },
   "outputs": [],
   "source": [
    "def euler_explicit_sir(N, T, y0, f, args):\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    y = np.zeros((N + 1, len(y0)))\n",
    "    y[0] = y0\n",
    "    h = T / N\n",
    "    for i in range(N):\n",
    "        y[i + 1] = y[i] + h * f(y[i], t[i], *args)\n",
    "    return t, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ca061c36b09f7200",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:46.330864Z",
     "start_time": "2025-03-25T10:54:46.325559Z"
    }
   },
   "outputs": [],
   "source": [
    "def heun_sir(N, T, y0, f, args):\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    y = np.zeros((N + 1, len(y0)))\n",
    "    y[0] = y0\n",
    "    h = T / N\n",
    "    for i in range(N):\n",
    "        k1 = h * f(y[i], t[i], *args)\n",
    "        k2 = h * f(y[i] + k1, t[i + 1], *args)\n",
    "        y[i + 1] = y[i] + 0.5 * (k1 + k2)\n",
    "    return t, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "afc8571b14a6b464",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:46.819777Z",
     "start_time": "2025-03-25T10:54:46.815301Z"
    }
   },
   "outputs": [],
   "source": [
    "def runge_kutta_order_4_sir(N, T, y0, f, args):\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    y = np.zeros((N + 1, len(y0)))\n",
    "    y[0] = y0\n",
    "    h = T / N\n",
    "    for i in range(N):\n",
    "        k1 = h * f(y[i], t[i], *args)\n",
    "        k2 = h * f(y[i] + 0.5 * k1, t[i] + 0.5 * h, *args)\n",
    "        k3 = h * f(y[i] + 0.5 * k2, t[i] + 0.5 * h, *args)\n",
    "        k4 = h * f(y[i] + k3, t[i + 1], *args)\n",
    "        y[i + 1] = y[i] + (k1 + 2 * k2 + 2 * k3 + k4) / 6\n",
    "    return t, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3ff4bbd4286d2b69",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:47.675843Z",
     "start_time": "2025-03-25T10:54:47.573186Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14796 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "t0, T = 52, 60\n",
    "N = 150\n",
    "\n",
    "y0 = np.array([10e6, 10, 0], dtype=np.float64)  # Initial conditions: S, I, R\n",
    "y0 = y0 / np.sum(y0)  # Normalize initial conditions\n",
    "a = 0.33  # Recovery rate\n",
    "r = 0.5  # Infection rate\n",
    "t = np.linspace(t0, T, N)\n",
    "\n",
    "sol = odeint(F, y0, t, args=(a, r))\n",
    "\n",
    "plt.plot(t, sol[:, 0], label=\"S(t)\")\n",
    "plt.plot(t, sol[:, 1], label=\"I(t)\")\n",
    "plt.plot(t, sol[:, 2], label=\"R(t)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c147c62e0a55bf28",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:48.871101Z",
     "start_time": "2025-03-25T10:54:48.363979Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36560 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "errors_EE = []\n",
    "errors_H = []\n",
    "errors_RK4 = []\n",
    "\n",
    "h_list = [0.05, 0.01, 0.1, 0.5, 1, 2, 4]\n",
    "\n",
    "\n",
    "def calculate_errors(sol_exact, sol_appr):\n",
    "    return np.max(\n",
    "        np.power(np.abs(sol_appr - sol_exact), 2)[\n",
    "            np.isfinite(np.power(np.abs(sol_appr - sol_exact), 2))\n",
    "        ]\n",
    "    )\n",
    "\n",
    "\n",
    "for h in h_list:\n",
    "    N = int(T / h)\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "\n",
    "    _, sol_appr_EE = euler_explicit_sir(N, T, y0, F, args=(a, r))\n",
    "    _, sol_appr_H = heun_sir(N, T, y0, F, args=(a, r))\n",
    "    _, sol_appr_RK4 = runge_kutta_order_4_sir(N, T, y0, F, args=(a, r))\n",
    "    sol_exact = odeint(F, y0, t, args=(a, r))\n",
    "\n",
    "    errors_EE.append(calculate_errors(sol_exact, sol_appr_EE))\n",
    "    errors_H.append(calculate_errors(sol_exact, sol_appr_H))\n",
    "    errors_RK4.append(calculate_errors(sol_exact, sol_appr_RK4))\n",
    "\n",
    "plt.plot(errors_EE, label=\"Euler Explicit\")\n",
    "plt.plot(errors_H, label=\"Heun\")\n",
    "plt.plot(errors_RK4, label=\"Runge Kutta order 4\")\n",
    "plt.yscale(\"log\")\n",
    "plt.xscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee44f24a438fb7b6",
   "metadata": {},
   "source": [
    "## Exercise 2.20 : Stability of Euler Explicit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4484c71f50b8e691",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:51.471808Z",
     "start_time": "2025-03-25T10:54:51.467799Z"
    }
   },
   "outputs": [],
   "source": [
    "def euler_implicit(N, T, u0, lam, type=int):\n",
    "    h = T / N\n",
    "    t = np.linspace(0, T, N + 1)\n",
    "    u = np.zeros(N + 1, dtype=type)\n",
    "    u[0] = u0\n",
    "    for n in range(N):\n",
    "        u[n + 1] = u[n] / (1 - lam * h)\n",
    "    return t, u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1b8a3146e7060c52",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-25T10:54:52.163532Z",
     "start_time": "2025-03-25T10:54:51.982495Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88148 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f(t, x, lam):\n",
    "    return x * lam\n",
    "\n",
    "\n",
    "T = 10 * 2 * np.pi\n",
    "lam = 1j\n",
    "N = 400\n",
    "x0 = 1\n",
    "\n",
    "t, x_appr_EE = euler_explicit(N, T, x0, f, type=complex, args=(lam,))\n",
    "_, x_appr_EI = euler_implicit(N, T, x0, lam, type=complex)\n",
    "x_exact = x0 * np.exp(lam * t)\n",
    "\n",
    "# Plot the real parts\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(t, np.real(x_appr_EI), label=\"Numerical Solution by EI\")\n",
    "plt.plot(t, np.real(x_appr_EE), label=\"Numerical Solution by EE\")\n",
    "plt.plot(t, np.real(x_exact), label=\"Exact Solution\", linestyle=\"--\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Real Part\")\n",
    "plt.legend()\n",
    "plt.title(\"Real Part of the Solution\")\n",
    "\n",
    "# Plot the imaginary parts\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(t, np.imag(x_appr_EI), label=\"Numerical Solution by EI\")\n",
    "plt.plot(t, np.imag(x_appr_EE), label=\"Numerical Solution by EE\")\n",
    "plt.plot(t, np.imag(x_exact), label=\"Exact Solution\", linestyle=\"--\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Imaginary Part\")\n",
    "plt.legend()\n",
    "plt.title(\"Imaginary Part of the Solution\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e770cb3923a39dcf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
