{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Case study: Spider-Man\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import functions from the modsim.py module\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "I'll start by getting the units we'll need from Pint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = UNITS.meter\n",
    "s = UNITS.second\n",
    "kg = UNITS.kilogram\n",
    "N = UNITS.newton\n",
    "degree = UNITS.degree\n",
    "radian = UNITS.radian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spider-Man"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case study we'll develop a model of Spider-Man swinging from a springy cable of webbing attached to the top of the Empire State Building.  Initially, Spider-Man is at the top of a nearby building, as shown in this diagram.\n",
    "\n",
    "![](diagrams/spiderman.png)\n",
    "\n",
    "The origin, `O⃗`, is at the base of the Empire State Building.  The vector `H⃗` represents the position where the webbing is attached to the building, relative to `O⃗`.  The vector `P⃗` is the position of Spider-Man relative to `O⃗`.  And `L⃗` is the vector from the attachment point to Spider-Man.\n",
    "\n",
    "By following the arrows from `O⃗`, along `H⃗`, and along `L⃗`, we can see that \n",
    "\n",
    "`H⃗ + L⃗ = P⃗`\n",
    "\n",
    "So we can compute `L⃗` like this:\n",
    "\n",
    "`L⃗ = P⃗ - H⃗`\n",
    "\n",
    "The goals of this case study are:\n",
    "\n",
    "1. Implement a model of this scenario to predict Spider-Man's trajectory.\n",
    "\n",
    "2. Choose the right time for Spider-Man to let go of the webbing in order to maximize the distance he travels before landing.\n",
    "\n",
    "3. Choose the best angle for Spider-Man to jump off the building, and let go of the webbing, to maximize range."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I'll create a `Params` object to contain the quantities we'll need:\n",
    "\n",
    "1. According to [the Spider-Man Wiki](http://spiderman.wikia.com/wiki/Peter_Parker_%28Earth-616%29), Spider-Man weighs 76 kg.\n",
    "\n",
    "2. Let's assume his terminal velocity is 60 m/s.\n",
    "\n",
    "3. The length of the web is 100 m.\n",
    "\n",
    "4. The initial angle of the web is 45 degrees to the left of straight down.\n",
    "\n",
    "5. The spring constant of the web is 40 N / m when the cord is stretched, and 0 when it's compressed.\n",
    "\n",
    "Here's a `Params` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = Params(height = 381 * m,\n",
    "                g = 9.8 * m/s**2,\n",
    "                mass = 75 * kg,\n",
    "                area = 1 * m**2,\n",
    "                rho = 1.2 * kg/m**3,\n",
    "                v_term = 60 * m / s,\n",
    "                length = 100 * m,\n",
    "                angle = (270 - 45) * degree,\n",
    "                k = 40 * N / m,\n",
    "                t_0 = 0 * s,\n",
    "                t_end = 30 * s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the initial position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initial_condition(params):\n",
    "    \"\"\"Compute the initial position and velocity.\n",
    "    \n",
    "    params: Params object\n",
    "    \"\"\"\n",
    "    height, length, angle = params.height, params.length, params.angle\n",
    "    \n",
    "    H⃗ = Vector(0, height)\n",
    "    theta = angle.to(radian)\n",
    "    x, y = pol2cart(theta, length)\n",
    "    L⃗ = Vector(x, y)\n",
    "    P⃗ = H⃗ + L⃗\n",
    "    V⃗ = Vector(0, 0) * m/s\n",
    "    \n",
    "    return State(P⃗=P⃗, V⃗=V⃗)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_condition(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now here's a version of `make_system` that takes a `Params` object as a parameter.\n",
    "\n",
    "`make_system` uses the given value of `v_term` to compute the drag coefficient `C_d`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_system(params):\n",
    "    \"\"\"Makes a System object for the given conditions.\n",
    "    \n",
    "    params: Params object\n",
    "    \n",
    "    returns: System object\n",
    "    \"\"\"\n",
    "    init = initial_condition(params)\n",
    "    \n",
    "    mass, g = params.mass, params.g\n",
    "    rho, area, v_term = params.rho, params.area, params.v_term\n",
    "    C_d = 2 * mass * g / (rho * area * v_term**2)\n",
    "    \n",
    "    return System(params, init=init, C_d=C_d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a `System`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = make_system(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drag and spring forces\n",
    "\n",
    "Here's drag force, as we saw in Chapter 22."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def drag_force(V⃗, system):\n",
    "    \"\"\"Compute drag force.\n",
    "    \n",
    "    V⃗: velocity Vector\n",
    "    system: `System` object\n",
    "    \n",
    "    returns: force Vector\n",
    "    \"\"\"\n",
    "    rho, C_d, area = system.rho, system.C_d, system.area\n",
    "    \n",
    "    mag = rho * V⃗.mag**2 * C_d * area / 2\n",
    "    direction = -V⃗.hat()\n",
    "    f_drag = direction * mag\n",
    "    return f_drag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "V⃗_test = Vector(10, 10) * m/s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "drag_force(V⃗_test, system)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's the 2-D version of spring force.  We saw the 1-D version in Chapter 21."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spring_force(L⃗, system):\n",
    "    \"\"\"Compute drag force.\n",
    "    \n",
    "    L⃗: Vector representing the webbing\n",
    "    system: System object\n",
    "    \n",
    "    returns: force Vector\n",
    "    \"\"\"\n",
    "    extension = L⃗.mag - system.length\n",
    "    if magnitude(extension) < 0:\n",
    "        mag = 0\n",
    "    else:\n",
    "        mag = system.k * extension\n",
    "        \n",
    "    direction = -L⃗.hat()\n",
    "    f_spring = direction * mag\n",
    "    return f_spring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "L⃗_test = Vector(0, -system.length-1*m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_spring = spring_force(L⃗_test, system)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the slope function, including acceleration due to gravity, drag, and the spring force of the webbing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def slope_func(state, t, system):\n",
    "    \"\"\"Computes derivatives of the state variables.\n",
    "    \n",
    "    state: State (x, y, x velocity, y velocity)\n",
    "    t: time\n",
    "    system: System object with g, rho, C_d, area, mass\n",
    "    \n",
    "    returns: sequence (vx, vy, ax, ay)\n",
    "    \"\"\"\n",
    "    P⃗, V⃗ = state\n",
    "    g, mass = system.g, system.mass\n",
    "    \n",
    "    H⃗ = Vector(0, system.height)\n",
    "    L⃗ = P⃗ - H⃗\n",
    "    \n",
    "    a_grav = Vector(0, -g)\n",
    "    a_spring = spring_force(L⃗, system) / mass\n",
    "    a_drag = drag_force(V⃗, system) / mass\n",
    "    \n",
    "    A⃗ = a_grav + a_drag + a_spring\n",
    "    \n",
    "    return V⃗, A⃗"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, let's test the slope function with the initial conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "slope_func(system.init, 0, system)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then run the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "results, details = run_ode_solver(system, slope_func)\n",
    "details"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the results\n",
    "\n",
    "We can extract the x and y components as `Series` objects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simplest way to visualize the results is to plot x and y as functions of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_position(P⃗):\n",
    "    x = P⃗.extract('x')\n",
    "    y = P⃗.extract('y')\n",
    "    plot(x, label='x')\n",
    "    plot(y, label='y')\n",
    "\n",
    "    decorate(xlabel='Time (s)',\n",
    "             ylabel='Position (m)')\n",
    "    \n",
    "plot_position(results.P⃗)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the velocities the same way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_velocity(V⃗):\n",
    "    vx = V⃗.extract('x')\n",
    "    vy = V⃗.extract('y')\n",
    "    plot(vx, label='vx')\n",
    "    plot(vy, label='vy')\n",
    "\n",
    "    decorate(xlabel='Time (s)',\n",
    "             ylabel='Velocity (m/s)')\n",
    "    \n",
    "plot_velocity(results.V⃗)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another way to visualize the results is to plot y versus x.  The result is the trajectory through the plane of motion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_trajectory(P⃗, **options):\n",
    "    x = P⃗.extract('x')\n",
    "    y = P⃗.extract('y')\n",
    "    plot(x, y, **options)\n",
    "    \n",
    "    decorate(xlabel='x position (m)',\n",
    "             ylabel='y position (m)')\n",
    "    \n",
    "plot_trajectory(results.P⃗, label='trajectory')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Letting go\n",
    "\n",
    "Now let's find the optimal time for Spider-Man to let go.  We have to run the simulation in two phases because the spring force changes abruptly when Spider-Man lets go, so we can't integrate through it.\n",
    "\n",
    "Here are the parameters for Phase 1, running for 9 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "params1 = Params(params, t_end=9*s)\n",
    "system1 = make_system(params1)\n",
    "results1, details1 = run_ode_solver(system1, slope_func)\n",
    "plot_trajectory(results1.P⃗, label='Phase 1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final conditions from Phase 1 are the initial conditions for Phase 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_final = get_last_label(results1) * s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the position Vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "init = results1.last_row()\n",
    "init.P⃗"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the velocity Vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "init.V⃗"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the `System` for Phase 2.  We can turn off the spring force by setting `k=0`, so we don't have to write a new slope function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "system2 = System(system1, t_0=t_final, t_end=t_final+10*s, init=init, k=0*N/m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an event function that stops the simulation when Spider-Man reaches the ground."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def event_func(state, t, system):\n",
    "    \"\"\"Stops when y=0.\n",
    "    \n",
    "    state: State object\n",
    "    t: time\n",
    "    system: System object\n",
    "    \n",
    "    returns: height\n",
    "    \"\"\"\n",
    "    P⃗, V⃗ = state\n",
    "    return P⃗.y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run Phase 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "results2, details2 = run_ode_solver(system2, slope_func, events=event_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_trajectory(results1.P⃗, label='Phase 1')\n",
    "plot_trajectory(results2.P⃗, label='Phase 2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can gather all that into a function that takes `t_release` and `V_0`, runs both phases, and returns the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_two_phase(t_release, V⃗_0, params):\n",
    "    \"\"\"Run both phases.\n",
    "    \n",
    "    t_release: time when Spider-Man lets go of the webbing\n",
    "    V_0: initial velocity\n",
    "    \"\"\"\n",
    "    params1 = Params(params, t_end=t_release, V⃗_0=V⃗_0)\n",
    "    system1 = make_system(params1)\n",
    "    results1, details1 = run_ode_solver(system1, slope_func)\n",
    "\n",
    "    t_0 = get_last_label(results1) * s\n",
    "    t_end = t_0 + 10 * s\n",
    "    init = results1.last_row()\n",
    "\n",
    "    system2 = System(system1, t_0=t_0, t_end=t_end, init=init, k=0*N/m)\n",
    "    results2, details2 = run_ode_solver(system2, slope_func, events=event_func)\n",
    "\n",
    "    results = results1.combine_first(results2)\n",
    "    return TimeFrame(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's a test run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_release = 9 * s\n",
    "V⃗_0 = Vector(0, 0) * m/s\n",
    "\n",
    "results = run_two_phase(t_release, V⃗_0, params)\n",
    "plot_trajectory(results.P⃗)\n",
    "x_final = results.P⃗.last_value().x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Animation\n",
    "\n",
    "Here's a draw function we can use to animate the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = results.P⃗.extract('x')\n",
    "ys = results.P⃗.extract('y')\n",
    "\n",
    "def draw_func(state, t):\n",
    "    set_xlim(xs)\n",
    "    set_ylim(ys)\n",
    "    x, y = state.P⃗\n",
    "    plot(x, y, 'bo')\n",
    "    decorate(xlabel='x position (m)',\n",
    "             ylabel='y position (m)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate(results, draw_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Maximizing range\n",
    "\n",
    "To find the best value of `t_release`, we need a function that takes possible values, runs the simulation, and returns the range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def range_func(t_release, params):\n",
    "    \"\"\"Compute the final value of x.\n",
    "    \n",
    "    t_release: time to release web\n",
    "    params: Params object\n",
    "    \"\"\"\n",
    "    V_0 = Vector(0, 0) * m/s\n",
    "    results = run_two_phase(t_release, V_0, params)\n",
    "    x_final = results.P⃗.last_value().x\n",
    "    print(t_release, x_final)\n",
    "    return x_final"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can test it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "range_func(9*s, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And run it for a few values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for t_release in linrange(3, 15, 3) * s:\n",
    "    range_func(t_release, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use `maximize_scalar` to find the optimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bounds = [6, 12] * s\n",
    "res = maximize_golden(range_func, bounds, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can run the simulation with the optimal value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_time = res.x\n",
    "V⃗_0 = Vector(0, 0) * m/s\n",
    "results = run_two_phase(best_time, V⃗_0, params)\n",
    "plot_trajectory(results.P⃗)\n",
    "x_final = results.P⃗.last_value().x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}