{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "import numpy as np\n",
    "from scipy.integrate import odeint\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "plt.rcParams[\"animation.html\"] = \"jshtml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = 1\n",
    "m2 = 1\n",
    "c1 = 0.1\n",
    "c2 = 0.1\n",
    "c3 = 0.1\n",
    "a1 = 0.1\n",
    "a2 = 0.1\n",
    "a3 = 0.1\n",
    "k1 = 0.05\n",
    "k2 = 0.21\n",
    "k3 = 0.11\n",
    "C1 = 0.15\n",
    "C2 = 0.30\n",
    "omega = 0.1\n",
    "\n",
    "f1 = lambda t: C1*np.cos(omega*t)\n",
    "f2 = lambda t: C2*np.cos(omega*t)\n",
    "\n",
    "def deriv(X, t):\n",
    "    \"\"\"Return the derivatives dx/dt and d2x/dt2.\"\"\"\n",
    "    x1, x2, xd1, xd2 = X\n",
    "    F1 = f1(t)\n",
    "    F2 = f2(t)\n",
    "    xdd1 = F1 - xd1*(c1 + c2) + xd2*c2 - x1*(k1 + k2) + x2*k2 - a1*x1**3 + a2*(x2 - x1)**3\n",
    "    xdd2 = F2 - xd2*(c2 + c3) + xd1*c1 - x2*(k2 + k3) + x1*k2 - a3*x2**3 + a2*(x2 - x1)**3\n",
    "    return xd1, xd2, xdd1/m1, xdd2/m2\n",
    "\n",
    "\n",
    "def solve_duffing(tmax, dt_per_period, t_trans, x0, v0):\n",
    "    \"\"\"Solve the Duffing equation with the standard odeint solver.\n",
    "    \n",
    "    Find the numerical solution to the Duffing equation using a suitable\n",
    "    time grid: tmax is the maximum time (s) to integrate to; t_trans is\n",
    "    the initial time period of transient behaviour until the solution\n",
    "    settles down (if it does) to some kind of periodic motion (these data\n",
    "    points are dropped) and dt_per_period is the number of time samples\n",
    "    (of duration dt) to include per period of the driving motion (frequency\n",
    "    omega).\n",
    "    \n",
    "    Returns the time grid, t (after t_trans), position, x, and velocity,\n",
    "    xdot, dt, and step, the number of array points per period of the driving\n",
    "    motion.\n",
    "    \n",
    "    \"\"\"\n",
    "    # Time point spacings and the time grid\n",
    "\n",
    "    period = 2*np.pi/omega\n",
    "    dt = 2*np.pi/omega / dt_per_period\n",
    "    step = int(period / dt)\n",
    "    t = np.arange(0, tmax, dt)\n",
    "    # Initial conditions: x, xdot\n",
    "    X0 = x0 + v0\n",
    "    X = odeint(deriv, X0, t)\n",
    "    idx = int(t_trans / dt)\n",
    "    return t[idx:], X[idx:], dt, step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the motion for a oscillator with initial positions and velocities\n",
    "x0, v0 = (0.5, -0.5), (0.5, -0.1)\n",
    "tmax, t_trans = 150, 0\n",
    "dt_per_period = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve the equation of motion.\n",
    "t, X, dt, pstep = solve_duffing(tmax, dt_per_period, t_trans, x0, v0)\n",
    "print(\"# samples: %d\" % len(t))\n",
    "print(\"dt: %.4f\" % dt)\n",
    "print(\"Steps per period: %d\" % pstep)\n",
    "x1, x2, xd1, xd2 = X.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ((ax11, ax12), (ax21, ax22)) = plt.subplots(nrows=2, ncols=2)\n",
    "\n",
    "ax11.plot(t, x1)\n",
    "ax11.set_xlabel(r'$x_1$')\n",
    "ax12.plot(t, x2)\n",
    "ax12.set_xlabel(r'$x_2$')\n",
    "ax21.plot(t, xd1)\n",
    "ax21.set_xlabel(r'$v_1$')\n",
    "ax22.plot(t, xd2)\n",
    "ax22.set_xlabel(r'$v_2$')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "fig2, ((ax11, ax12), (ax21, ax22)) = plt.subplots(nrows=2, ncols=2)\n",
    "\n",
    "# Positions\n",
    "ax11.set_xlabel(r'$x_1$')\n",
    "ax11.set_ylabel(r'$x_2$')\n",
    "ln11, = ax11.plot([], [])\n",
    "ax11.set_xlim([np.min(x1), np.max(x1)])\n",
    "ax11.set_ylim([np.min(x2), np.max(x2)])\n",
    "\n",
    "# Velocities\n",
    "ax12.set_xlabel(r'$v_1$')\n",
    "ax12.set_ylabel(r'$v_2$')\n",
    "ln12, = ax12.plot([], [])\n",
    "ax12.set_xlim([np.min(xd1), np.max(xd1)])\n",
    "ax12.set_ylim([np.min(xd2), np.max(xd2)])\n",
    "\n",
    "# Phase spaces\n",
    "ax21.set_xlabel(r'$\\mathbf{x}$')\n",
    "ax21.set_ylabel(r'$\\mathbf{v}$')\n",
    "ln21a, = ax21.plot([], [])\n",
    "ln21b, = ax21.plot([], [])\n",
    "ax21.set_xlim([min(np.min(x1), np.min(x2)), max(np.max(x1), np.max(x2))])\n",
    "ax21.set_ylim([min(np.min(xd1), np.min(xd2)), max(np.max(xd1), np.max(xd2))])\n",
    "\n",
    "# F(t)\n",
    "F1 = f1(t)\n",
    "F2 = f2(t)\n",
    "ax22.set_xlabel(r'$t$')\n",
    "ax22.set_ylabel(r'$\\mathbf{F}(t)$')\n",
    "ln22a, = ax22.plot([], [])\n",
    "ln22b, = ax22.plot([], [])\n",
    "ax22.set_xlim([t[0], t[-1]])\n",
    "ax21.set_ylim([min(np.min(F1), np.min(F2)), max(np.max(F1), np.max(F2))])\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "def animate(i):\n",
    "    ln11.set_data(x1[:i], x2[:i])\n",
    "    ln12.set_data(xd1[:i], xd2[:i])\n",
    "    ln21a.set_data(x1[:i], xd1[:i])\n",
    "    ln21b.set_data(x2[:i], xd2[:i])\n",
    "    ln22a.set_data(t[:i], F1[:i])\n",
    "    ln22b.set_data(t[:i], F2[:i])\n",
    "    \n",
    "\n",
    "anim = animation.FuncAnimation(fig2, animate, frames=len(t), interval=100, repeat_delay=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
