interactivePlot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "215335dc1ef744d4b5501ce059dceb8e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.5, description='Beta 0:', max=1.0, step=0.01), FloatSlider(value=0.2…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.integrate import solve_ivp\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import interactive, HBox, VBox\n",
    "\n",
    "# entire sir model\n",
    "def sir_model(t, y, beta_0, beta_1, gamma_0, alpha, N, L, k, t_0):\n",
    "    S, I, R = y # each value will be plotted against time\n",
    "    dSdt = -beta_0 * (1 + beta_1 * np.cos(2 * np.pi * t)) * S * I / N * L / (1 + np.exp(-k * (t - t_0))) # susceptible\n",
    "    dIdt = beta_0 * (1 + beta_1 * np.cos(2 * np.pi * t)) * S * I / N - gamma_0 * (1 - alpha * I / N) * I # infected\n",
    "    dRdt = gamma_0 * (1 - alpha * I / N) * I # recovered\n",
    "    return [dSdt, dIdt, dRdt]\n",
    "\n",
    "# function to compute the derivatives for the vector field\n",
    "def compute_derivatives(S, I, beta_0, beta_1, gamma_0, alpha, N, L, k, t_0, t): # vector field will be susceptible vs infected\n",
    "    dSdt = -beta_0 * (1 + beta_1 * np.cos(2 * np.pi * t)) * S * I / N * L / (1 + np.exp(-k * (t - t_0)))\n",
    "    dIdt = beta_0 * (1 + beta_1 * np.cos(2 * np.pi * t)) * S * I / N - gamma_0 * (1 - alpha * I / N) * I\n",
    "    return dSdt, dIdt\n",
    "\n",
    "# plotting function\n",
    "def plot_sir_and_vector_field(beta_0, beta_1, gamma_0, alpha, N, L, k, t_0, S0, I0, R0, t_end):\n",
    "    # solve the SIR model\n",
    "    y0 = [S0, I0, R0]\n",
    "    t_span = [0, t_end]\n",
    "    sol = solve_ivp(sir_model, t_span, y0, args=(beta_0, beta_1, gamma_0, alpha, N, L, k, t_0), t_eval=np.linspace(0, t_end, 1000))\n",
    "\n",
    "    # plotting the model\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(20, 6))\n",
    "    \n",
    "    ax[0].plot(sol.t, sol.y[0], label='Susceptible')\n",
    "    ax[0].plot(sol.t, sol.y[1], label='Infected')\n",
    "    ax[0].plot(sol.t, sol.y[2], label='Recovered')\n",
    "    ax[0].set_xlabel('Time')\n",
    "    ax[0].set_ylabel('Population')\n",
    "    ax[0].set_title('SIR Model')\n",
    "    ax[0].legend()\n",
    "    ax[0].grid(True)\n",
    "\n",
    "    # creating vector field\n",
    "    S_range = np.linspace(0, N, 20)\n",
    "    I_range = np.linspace(0, N, 20)\n",
    "    S_grid, I_grid = np.meshgrid(S_range, I_range)\n",
    "    dSdt_grid, dIdt_grid = compute_derivatives(S_grid, I_grid, beta_0, beta_1, gamma_0, alpha, N, L, k, t_0, 0)\n",
    "\n",
    "    # plotting the vector field\n",
    "    ax[1].quiver(S_grid, I_grid, dSdt_grid, dIdt_grid)\n",
    "    ax[1].set_xlabel('Susceptible Population (S)')\n",
    "    ax[1].set_ylabel('Infected Population (I)')\n",
    "    ax[1].set_title('Vector Field for S vs I')\n",
    "    ax[1].axis([0, N, 0, N])\n",
    "    ax[1].grid(True)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# since i find it annoying to change each value, im gonna setup sliders for this\n",
    "beta_0_slider = widgets.FloatSlider(value=0.5, min=0, max=1, step=0.01, description='Beta 0:')\n",
    "beta_1_slider = widgets.FloatSlider(value=0.2, min=0, max=1, step=0.01, description='Beta 1:')\n",
    "gamma_0_slider = widgets.FloatSlider(value=0.1, min=0, max=1, step=0.01, description='Gamma 0:')\n",
    "alpha_slider = widgets.FloatSlider(value=0.1, min=0, max=1, step=0.01, description='Alpha:')\n",
    "N_slider = widgets.FloatSlider(value=16000, min=1, max=100000, step=100, description='Total Population (N):')\n",
    "L_slider = widgets.FloatSlider(value=0.9, min=0, max=1, step=0.01, description='L:')\n",
    "k_slider = widgets.FloatSlider(value=0.2, min=0, max=1, step=0.01, description='k:')\n",
    "t_0_slider = widgets.FloatSlider(value=5, min=0, max=10, step=0.1, description='t_0:')\n",
    "S0_slider = widgets.FloatSlider(value=9900, min=0, max=10000, step=100, description='S0:')\n",
    "I0_slider = widgets.FloatSlider(value=10, min=0, max=1000, step=10, description='I0:')\n",
    "R0_slider = widgets.FloatSlider(value=0, min=0, max=1000, step=10, description='R0:')\n",
    "t_end_slider = widgets.FloatSlider(value=100, min=0, max=1000, step=10, description='End Time:')\n",
    "\n",
    "# setup interactive plot\n",
    "interactive_plot = interactive(plot_sir_and_vector_field, beta_0=beta_0_slider, beta_1=beta_1_slider, gamma_0=gamma_0_slider,\n",
    "                               alpha=alpha_slider, N=N_slider, L=L_slider, k=k_slider, t_0=t_0_slider,\n",
    "                               S0=S0_slider, I0=I0_slider, R0=R0_slider, t_end=t_end_slider)\n",
    "\n",
    "display(interactive_plot) # displays the plot; notebook only command :(\n"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}