{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('default')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Travel Time for Seismic Waves\n",
    "## Lecture 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Review solving ODE numerically"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider the intial value problem of solving the following ordinary differential equation:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{dy}{dt} = f(y), \\quad y(0) = y_0\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In such problems, $y(t)$ is an unknown function that needs to be determined.  The right hand side of the differential equation, $f(y)$, represents a known function of $y$ that needs to be given explicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have seen a number of methods (Euler, Midpoint, RK4) that can be used to numerically solve such problems.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Projectile Motion\n",
    "\n",
    "At its simplest, projectile motion involves solving\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{d^2 x}{dt^2} &= 0 \\\\\n",
    "\\frac{d^2 y}{dt^2} &= -g\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "given the initial conditions for $x(0)$, $y(0)$, and the initial velocities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have learned that to solve such problems we can rewrite the equations as a system of four first order differential equations:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align}\n",
    "\\frac{d x}{dt} &= v_x \\\\\n",
    "\\frac{d y}{dt} &= v_y \\\\\n",
    "\\frac{d v_x}{dt} &= 0 \\\\\n",
    "\\frac{d v_y}{dt} &= -g \\\\ \\,\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "given the initial conditions $x(0) = x_0$, $y(0) = y_0$, $v_x(0) = v_{x0}$, $v_y(0) = v_{y0}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numerical solution\n",
    "\n",
    "Let's solve the projectile motion problem using the midpoint scheme."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Midpoint scheme (review)\n",
    "\n",
    "1\\. Estimate slope $s_1$ at $t$\n",
    "\n",
    "$$ s_1 = f(y(t), t) $$\n",
    "\n",
    "2\\. Use $s_1$ to estimate the midpoint between $t$ and $t + \\Delta t$\n",
    "\n",
    "\\begin{align}\n",
    "y^* &= y ( t + \\Delta t) \\\\\n",
    "       &= y(t) + \\frac{\\Delta t}{2} s_1\n",
    "\\end{align}\n",
    " \t\n",
    "\n",
    "3\\. Use $y^*$ to get the the slope $s_2$ at the midpoint\n",
    "\n",
    "$$ s_2 = f( y^*, t +  \\frac{\\Delta t}{2} ) $$\n",
    "\n",
    "4\\. Use $s_2$ to estimate $y(t + \\Delta t)$\n",
    "\n",
    "$$ y(t + \\Delta t) = y(t) + \\Delta t s_2 $$\n",
    "\n",
    "### Midpoint scheme algorithm\n",
    "\n",
    "\\begin{align}\n",
    "s_1 &= f(y_i, t_i) \\\\\n",
    "y^* &= y_i + \\Delta t / 2 s_1 \\\\\n",
    "s_2 &= f(y^*, t_i + \\Delta t /2) \\\\\n",
    "y_{i+1} &= y_i + \\Delta t s_2\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_1(tmax = 2, dt = 0.01,\n",
    "            x0 = 0, y0 = 0, vx0 = 3, vy0 = 10):  \n",
    "    \n",
    "    g = 9.81\n",
    "    \n",
    "    # Allocate memory for arrays and set to zero\n",
    "    N = round(tmax/dt)\n",
    "    x = np.zeros(N)\n",
    "    y = np.zeros(N)\n",
    "    vx = np.zeros(N)\n",
    "    vy = np.zeros(N)\n",
    "    t = np.zeros(N)\n",
    "\n",
    "    # Initial values\n",
    "    x[0] = x0\n",
    "    y[0] = y0\n",
    "    vx[0] = vx0\n",
    "    vy[0] = vy0\n",
    "    t[0] = 0\n",
    "    \n",
    "    # Calculate the solution using the Midpoint Method:\n",
    "    for i in range(N-1):\n",
    "        # Estimate slope s1 at  t\n",
    "        s1_x = vx[i]\n",
    "        s1_y = vy[i]\n",
    "        s1_vx = 0\n",
    "        s1_vy = -g\n",
    "        \n",
    "        # Use s1 to estimate the midpoint between t and t+Δt\n",
    "        x_tmp  = x[i]  + s1_x*dt/2\n",
    "        y_tmp  = y[i]  + s1_y*dt/2\n",
    "        vx_tmp = vx[i] + s1_vx*dt/2\n",
    "        vy_tmp = vy[i] + s1_vy*dt/2\n",
    "        \n",
    "        # Estimate the slope s2 at the midpoint\n",
    "        s2_x = vx_tmp\n",
    "        s2_y = vy_tmp\n",
    "        s2_vx = 0\n",
    "        s2_vy = -g\n",
    "       \n",
    "        # Use s2 to estimate y(t+Δt)\n",
    "        x[i+1]  = x[i]  + s2_x*dt\n",
    "        y[i+1]  = y[i]  + s2_y*dt\n",
    "        vx[i+1] = vx[i] + s2_vx*dt\n",
    "        vy[i+1] = vy[i] + s2_vy*dt\n",
    "     \n",
    "        t[i+1]  = t[i] + dt\n",
    "\n",
    "    return x, y, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, t = solve_1()\n",
    "\n",
    "plt.plot(x, y)\n",
    "plt.xlabel('x (m)')\n",
    "plt.ylabel('y (m)')\n",
    "plt.title('Projectile Motion solved with Midpoint Scheme')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector formulation\n",
    "\n",
    "Going back to the general form of an ODE we can also think about the unknown function as being a *vector-valued function* such as\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{d \\vec{q} }{dt} = \\vec{F}( \\vec{q} ), \\quad \\vec{q}(0) = \\vec{q}_0 \n",
    "\\end{align}\n",
    "\n",
    "where $ \\vec{q}(t) = \\left( q_1(t), q_2(t), \\ldots, q_m(t) \\right)$ is a vector in which every component is a function. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "We can define a vector for the projectile motion problem like this\n",
    "\n",
    "$$\n",
    "\\vec{q}(t) = \\left( x(t), y(t), v_x(t), v_y(t) \\right)\n",
    "$$\n",
    "\n",
    "This vector represents the **state** of the projectile motion system. Likewise, we can write down \n",
    "\n",
    "\\begin{align}\n",
    "\\vec{F}( \\vec{q} ) &= \\vec{F}(x, y, v_x, v_y) = (v_x, v_y, 0, -g) \\\\\n",
    "\\vec{q}_0 &= \\left( x_0, y_0, v_{x0}, v_{y0} \\right) \\\\\n",
    "\\end{align}\n",
    "\n",
    "This is just a compact rewriting of our system of four first-order DEs in a vector formulation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This allows us to generalize our midpoint scheme as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solver(F, q0, t):  \n",
    "    \"\"\"\n",
    "    Solve the ODE dq/dt = F(q) with q(0) = q0\n",
    "    \n",
    "    where F is a vector valued function\n",
    "          q0 is a vector of initial conditions\n",
    "          t is time array variable\n",
    "    \"\"\"\n",
    "\n",
    "    # count number points of time to solve\n",
    "    N = len(t)\n",
    "    # count number of indepedent variables\n",
    "    m = len(q0)\n",
    "    # Allocate memory set arrays to zero\n",
    "    q = np.zeros((m, N))\n",
    "    \n",
    "    # Initial values\n",
    "    q[:, 0] = q0\n",
    "    \n",
    "    # Calculate the solution using the Midpoint Method:\n",
    "    for i in range(N-1):\n",
    "        dt = t[i+1] - t[i]\n",
    "        \n",
    "        # Estimate slope s1 at  t\n",
    "        s1 = F( q[:, i], t[i] )\n",
    "        s1 = np.asarray(s1) # ensure s1 is a numpy array\n",
    "        \n",
    "        # Use s1 to estimate the midpoint between t and t+Δt\n",
    "        q_tmp = q[:, i] + s1*dt/2\n",
    "        \n",
    "        # Estimate the slope s2 at the midpoint\n",
    "        s2 = F( q_tmp, t[i]+ dt/2 )\n",
    "        s2 = np.asarray(s2) # ensure s2 is a numpy array\n",
    "        \n",
    "        # Use s2 to estimate y(t+Δt)\n",
    "        q[:, i+1] = q[:, i] + s2*dt\n",
    "\n",
    "    return q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that our solver is now very general. The same code could be used for many different differential equations.  \n",
    "\n",
    "We can formulate our specific projectile motion problem by defining a **particular** right-hand-side function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = 9.81\n",
    "    \n",
    "tmax = 2\n",
    "dt = 0.01\n",
    "\n",
    "# create an array for the time variables\n",
    "t = np.arange(0, tmax, dt)\n",
    "\n",
    "# define the right hand side, F(q)\n",
    "def F( q, t ):\n",
    "    # separate out the variables\n",
    "    x, y, vx, vy = q\n",
    "    # evaluate the right hand side of the ODE\n",
    "    dqdt = [vx, vy, 0, -g ]\n",
    "    \n",
    "    return dqdt\n",
    "\n",
    "# define initial values (x0, y0, vx0, vy0)\n",
    "q0 = [0, 0, 3, 10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then all we need to do is pass in this specific function for F and the initial conditions `q0` into our general solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, vx, vy = solver(F, q0, t)\n",
    "\n",
    "plt.plot(x,y)\n",
    "plt.xlabel('x (m)')\n",
    "plt.ylabel('y (m)')\n",
    "plt.title('Projectile Motion solved with Midpoint Scheme')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This abstraction separating out the physical model from the numerical solver, also allows us to try new solvers with almost *no changes the code at all*.\n",
    "\n",
    "The Python package `scipy` comes with its own ODE solver.  Here is the same projectile motion problem using a different solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we need to import the SciPy integrate module\n",
    "from scipy import integrate\n",
    "\n",
    "sol = integrate.odeint(F, q0, t)\n",
    "x, y, vx, vy = sol.T\n",
    "\n",
    "plt.plot(x,y)\n",
    "plt.xlabel('x (m)')\n",
    "plt.ylabel('y (m)')\n",
    "plt.title('Projectile Motion solved with integrate.odeint')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise:\n",
    "> Compare the previous two cell blocks.  What differences can you observe?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aside: Introduction to Seismology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://www.youtube.com/watch?v=22m27MhzSQs\n",
    "\n",
    "http://rallen.berkeley.edu/teaching/F04_GEO594_IntroAppGeophys/Lectures/L13_WavesAndRaysII.pdf\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ray Tracing in Seismology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{align}\n",
    "\\frac{dx}{dt} &= v \\sin \\theta \\\\\n",
    "\\frac{dy}{dt} &= -v \\cos \\theta \\\\\n",
    "\\frac{d\\theta}{dt} &= - \\cos \\theta \\frac{\\partial v}{\\partial x} - \\sin \\theta \\frac{\\partial v}{\\partial y}  \\\\\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where\n",
    "\n",
    "$$ v = v(x, y) $$\n",
    "\n",
    "is the velocity of a seismic waves that depends on the physical composition of the rock within the ground.\n",
    "\n",
    "#### Forward problem\n",
    "$ v(x,y)$ is assumed to be known and the goal is to determine the path of the rays.\n",
    "\n",
    "#### Inverse problem\n",
    "Measurements of the travel times of the rays are known and the goal is to estimate $v(x,y)$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forward Problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a rock structure that is getting linearly denser with depth.\n",
    "\n",
    "#### Birch's Law\n",
    "\n",
    "$$ v = 0.33 \\rho + 2.20 $$\n",
    "\n",
    "were $v$ is in  km/s if  $\\rho$ is in $10^3$ km/m$^3$.\n",
    "\n",
    "So,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.arange(-40, 0, 0.1)\n",
    "x = np.arange(0, 120, 0.1)\n",
    "\n",
    "def ρ(x, y):\n",
    "    return (2000 - 1500/40*y)/1000 \n",
    "\n",
    "def v(x, y):\n",
    "    return 0.33 + 2.20*ρ(x, y)\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize=(10,6))\n",
    "plt.sca(axes[0])\n",
    "plt.plot(ρ(0, y), y)\n",
    "plt.xlabel(r'$\\rho$')\n",
    "plt.ylabel(r'$y$ (km)')\n",
    "plt.sca(axes[1])\n",
    "plt.plot(v(0, y), y)\n",
    "plt.xlabel(r'$v$ km/s')\n",
    "plt.ylabel(r'$y$ (km)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's send out a ray!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sin = np.sin\n",
    "cos = np.cos\n",
    "π = np.pi\n",
    "\n",
    "tmax = 22\n",
    "dt = 0.1\n",
    "\n",
    "# create an array for the time variables\n",
    "t = np.arange(0, tmax, dt)\n",
    "\n",
    "# define the right hand side, F(q)\n",
    "def F( q, t ):\n",
    "    # separate out the variables\n",
    "    x, y, θ = q\n",
    "    \n",
    "    # evaluate the right hand side of the ODE\n",
    "    dxdt = v(x,y)*sin(θ)\n",
    "    dydt = -v(x,y)*cos(θ)\n",
    "    dθdt = -cos(θ)*dvdx(x, y) - sin(θ)*dvdy(x,y)\n",
    "    dqdt = [dxdt, dydt, dθdt]\n",
    "    \n",
    "    return dqdt\n",
    "\n",
    "# define initial values (x0, y0, θ0)\n",
    "q0 = [0, 0, π/4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also need the derivatives of $v(x,y)$.  We can estimate those numerically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Centered difference scheme \n",
    "def NDc(f, x0, dx):\n",
    "    return (f(x0+dx) - f(x0-dx)) / (2*dx)\n",
    "\n",
    "def dvdx(x, y):\n",
    "    return NDc(lambda x: v(x,y), x, 1)\n",
    "\n",
    "def dvdy(x, y):\n",
    "    return NDc(lambda y: v(x,y), y, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = integrate.odeint(F, q0, t)\n",
    "x, y, θ = sol.T\n",
    "\n",
    "plt.plot(x, y)\n",
    "plt.xlabel('x (km)')\n",
    "plt.ylabel('y (km)')\n",
    "plt.title('Ray Trace')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could do this for a number of different initial angles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vary_theta(x0=0, y0=0, θmin = π/4, θmax = π/2):\n",
    "    for θ0 in np.arange(θmin, θmax, 0.1):\n",
    "        q0 = [x0, y0, θ0]\n",
    "\n",
    "        sol = integrate.odeint(F, q0, t)\n",
    "        x, y, θ = sol.T\n",
    "\n",
    "        plt.plot(x, y, 'k-')\n",
    "    plt.xlabel('x (km)')\n",
    "    plt.ylabel('y (km)')\n",
    "    plt.ylim(-30, 0)\n",
    "    plt.xlim(0, 120)\n",
    "    plt.title('Ray Traces')\n",
    "    \n",
    "vary_theta()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Probably a good idea to show the density/velocity as a background to these traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_velocity():\n",
    "    X, Y = np.mgrid[0:120:0.1, -30:0:0.1]\n",
    "    V = v(X,Y)\n",
    "    plt.contourf(X, Y, V, 10, cmap=plt.cm.Blues)\n",
    "    plt.xlabel('x (km)')\n",
    "    plt.ylabel('y (km)')\n",
    "    plt.colorbar()\n",
    "    \n",
    "plot_velocity()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course we could put these together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_velocity()\n",
    "vary_theta()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if there were some variations of rock density.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def ρ_perturbation(x, y):\n",
    "    r =0\n",
    "    r += 1*np.exp(-(x-60)**2/ 20 - (y+15)**2/20)\n",
    "    r += 0.1*np.exp(-(x-20)**2/ 30 - (y+5)**2/20)\n",
    "    \n",
    "    return r\n",
    "    \n",
    "def ρ(x, y):\n",
    "    return (2000 - 1500/40*y)/1000 + ρ_perturbation(x, y)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_velocity()\n",
    "vary_theta()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seismic pulse at depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_velocity()\n",
    "vary_theta(x0=80, y0 = -25, θmin = -3*π/4, θmax = -π/2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}