day2_solving_ODEs_exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](images/uspas.png)\n",
    "# VUV and X-ray Free Electron Lasers\n",
    "# Lab Day 2: Solving ODEs with Python\n",
    "##### Additional Python references/examples: \n",
    "- https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html\n",
    "- https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html\n",
    "\n",
    "##### Instructors: D. Nguyen, P. Anisimov, N. Neveu\n",
    "##### Teaching Assistant: Y.S. Li"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing needed packages\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.integrate import solve_ivp\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use(r'./PaperDoubleFig.mplstyle')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](images/pendulum.png)\n",
    "\n",
    "#### Lecture 1, ~slide 60 defines the ODE for pendulum and FEL energy-phase equations. \n",
    "\n",
    "## Pendulum:   \n",
    "\n",
    "To solve the second order DE with Python, we need to break it down into two first order ODEs: \n",
    "\n",
    "$\\frac{d\\theta}{dt} = \\omega$  \n",
    "\n",
    "$\\frac{d\\ddot\\theta}{dt} = \\frac{d\\omega}{dt}=-\\frac{g}{l}sin\\theta$  \n",
    "\n",
    "We can solve the ODE's with Python's solve_ivp function from scipy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pendulum equations Python\n",
    "def pend(t, y, c):\n",
    "    '''\n",
    "    y = array of theta, omega, c\n",
    "    t = time\n",
    "    c = g/l \n",
    "    '''\n",
    "    theta, omega = y\n",
    "    dydt = np.array([omega, -c*np.sin(theta)])\n",
    "    return dydt\n",
    "\n",
    "#initial conditions\n",
    "th = np.pi-0.1 # pendulum nearly vertical\n",
    "om = 0.0 # at rest\n",
    "c  = 5\n",
    "\n",
    "y0  = np.array([th,om])\n",
    "sol = solve_ivp(pend, [0,10], y0, args=(c,), max_step=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Take a look at the solve_ivp output\n",
    "sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot theta and omega vs. t\n",
    "plt.plot(sol.t, sol.y[0], 'b', label=r'$\\theta(t)$')\n",
    "plt.plot(sol.t, sol.y[1], 'g', label=r'$\\omega(t)$')\n",
    "plt.legend(loc='best')\n",
    "plt.xlabel('time')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initial conditions\n",
    "th = np.pi-0.1 # pendulum nearly vertical\n",
    "om = 2.0  \n",
    "c  = 5\n",
    "\n",
    "y0  = np.array([th,om])\n",
    "sol = solve_ivp(pend, [0,10], y0, args=(c,), max_step=0.1)\n",
    "\n",
    "plt.plot(sol.t, sol.y[0], 'b', label=r'$\\theta(t)$')\n",
    "plt.plot(sol.t, sol.y[1], 'g', label=r'$\\omega(t)$')\n",
    "plt.legend(loc='best')\n",
    "plt.xlabel('time')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Write a sentence explaining the behavior of theta in this plot. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Use solve_ivp to plot $\\psi(t)$ and $\\eta(t)$:\n",
    "## FEL energy-phase: \n",
    "\n",
    "$\\frac{d\\psi}{dz} = 2 k_u \\eta$  \n",
    "\n",
    "$\\frac{d\\eta}{dz} = - |a|sin\\psi$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fel(t, y, k, a):\n",
    "    '''\n",
    "    t - time\n",
    "    y - [psi, eta]\n",
    "    k - undulator parameter\n",
    "    '''\n",
    "    ly = len(y)\n",
    "    hy = int(ly/2)\n",
    "    psi  = y[0:hy]\n",
    "    eta  = y[hy:]\n",
    "#     dpdz = \n",
    "#     dedz = \n",
    "    dydz = np.concatenate([dpdz, dedz])\n",
    "    \n",
    "    return dydz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initial conditions\n",
    "ne = 100 #number of electrons\n",
    "a  = 1\n",
    "k  = 3.5 \n",
    "eta0 = np.random.uniform(-0.5,0.5,ne)\n",
    "psi0 = np.linspace(-np.pi,np.pi,ne)\n",
    "\n",
    "# solve ODE's\n",
    "y1   = np.concatenate([psi0, eta0])\n",
    "solf = solve_ivp(fel, [0,5], y1, args=(k,a), max_step=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot psi and eta vs. t\n",
    "plt.plot(solf.t, solf.y[0:ne,:].T)\n",
    "plt.xlabel('t')\n",
    "plt.ylabel(r'$\\psi(z)$')\n",
    "plt.grid()\n",
    "plt.show()\n",
    "\n",
    "plt.plot(solf.t, solf.y[ne:,:].T)\n",
    "plt.xlabel('t')\n",
    "plt.ylabel(r'$\\eta(z)$')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Change the constant 'a' to a time dependant function\n",
    "Give solve_ivp 'a' as a function of time. \n",
    "i.e. a = np.cos (w/o parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fel2(t, y, k, a):\n",
    "    '''\n",
    "    t - time\n",
    "    y - [psi, eta]\n",
    "    k - undulator parameter\n",
    "    '''\n",
    "\n",
    "    \n",
    "    return dydz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initial conditions\n",
    "ne = 100\n",
    "# at = \n",
    "k  = 3.5 \n",
    "eta0 = np.random.uniform(-0.5,0.5,ne)\n",
    "psi0 = np.linspace(-np.pi,np.pi,ne)\n",
    "\n",
    "y2    = np.concatenate([psi0, eta0])\n",
    "solf2 = solve_ivp(fel2, [0,5], y2, args=(k,at), max_step=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "py3"
  },
  "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}