{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D Travelling Signal (sine with linear dispersion)\n",
    "\n",
    "\n",
    "$$\n",
    "u(x, t) = \\sin\\left( \\frac{2 \\pi x}{\\lambda} -2\\pi f t\\right)\n",
    "\\\\\n",
    "= \\sin\\left( 2 \\pi f \\, ( \\frac{x}{v(f)} -t) \\right)\n",
    "$$\n",
    "\n",
    "$\\lambda = \\dfrac{v}{f}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.fft as ft\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import axes3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TravellingSine:\n",
    "    \"\"\"\n",
    "    Scalar Travelling Sine wave with a constant dispersion relation per frequency\n",
    "    \"\"\"\n",
    "    def __init__(self, x_0, t_0, frequencies):\n",
    "        self.x_0 = x_0\n",
    "        self.t_0 = t_0\n",
    "        self.freqs = np.array(frequencies).reshape(-1)\n",
    "\n",
    "    def __call__(self, x_f, t_f, phase_velocities = None):\n",
    "        x_f = np.array(x_f)\n",
    "        t_f = np.array(t_f)\n",
    "        return self._signal(self.x_0, x_f, self.t_0, t_f, self.freqs, phase_velocities)\n",
    "\n",
    "    def _signal(self, x_0, x_f, t_0, t_f, freqs, phase_velocities = None, debug = 0):\n",
    "\n",
    "        # broadcasting between x_f and t_f\n",
    "        if np.ndim(x_f) != 0 or np.ndim(t_f) != 0:\n",
    "            if np.ndim(x_f) == 0:\n",
    "                x_f = x_f * np.ones(t_f.shape)\n",
    "            if np.ndim(t_f) == 0:\n",
    "                t_f = t_f * np.ones(x_f.shape)\n",
    "        \n",
    "        assert len(x_f) == len(t_f), \"x_f and t_f MUST have equal length: {} and {}\".format(len(x_f), len(t_f))\n",
    "        \n",
    "        # broadcasting between frequencies and phase velocities\n",
    "        if phase_velocities is None:\n",
    "            phase_velocities = 1\n",
    "        \n",
    "        if callable(phase_velocities):\n",
    "            phase_velocities = phase_velocities(freqs)\n",
    "\n",
    "        if np.ndim(phase_velocities) == 0:\n",
    "            phase_velocities = phase_velocities * np.ones(self.freqs.shape)\n",
    "        \n",
    "        assert len(freqs) == len(phase_velocities), \"freqs and phase_velocities MUST have equal length: {} and {}\".format(len(freqs), len(phase_velocities))\n",
    "        \n",
    "        # get the wavelengths for each frequency\n",
    "        wavelengths = phase_velocities / freqs\n",
    "        \n",
    "        # make sure these are 1d arrays\n",
    "        time_diff  = np.array(t_0 - t_f).reshape(-1)        \n",
    "        space_diff = np.array(x_0 - x_f).reshape(-1)       \n",
    "    \n",
    "        # ndarray of phase components\n",
    "        time_phase  = np.outer(2*np.pi*freqs, time_diff)\n",
    "        space_phase = np.outer(2*np.pi/wavelengths, space_diff)\n",
    "        \n",
    "        if debug:\n",
    "            if debug > 1:\n",
    "                #print('time_diff', time_diff)\n",
    "                print('frequencies', freqs)\n",
    "                #print('time_phase', time_phase)\n",
    "\n",
    "                #print('space_diff', space_diff)\n",
    "                print('wavelengths', wavelengths)\n",
    "                #print('space_phase', space_phase)\n",
    "            \n",
    "            print('time_phase shape', time_phase.shape)\n",
    "            print('space_phase shape', space_phase.shape)\n",
    "    \n",
    "        signal = np.sin( time_phase + space_phase )\n",
    "\n",
    "        if freqs.shape:\n",
    "            signal = signal.sum(axis=0)\n",
    " \n",
    "        return signal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85064 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "frequencies = 1/np.array([2,-2])\n",
    "velocities = np.array([1, 4])\n",
    "signal  = TravellingSine(0, 0, frequencies)\n",
    "\n",
    "zeros = np.zeros(500)\n",
    "x = np.linspace(-10, 10, 500)\n",
    "t = x\n",
    "\n",
    "x_mesh, t_mesh = np.meshgrid(x,t, sparse=True)\n",
    "\n",
    "fig = plt.figure()\n",
    "plt.suptitle(\"Travelling sine wave\")\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot(x, zeros, signal(x, zeros, velocities), label='right moving')\n",
    "ax.plot(x, t, signal(x, t, velocities), label='left moving')\n",
    "ax.plot(x, t, signal(x, t, 1), label='flat line')\n",
    "ax.set_xlabel(\"spatial distance\")\n",
    "ax.set_ylabel('temporal distance')\n",
    "ax.set_zlabel(\"amplitude\")\n",
    "ax.legend()\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# animation\n",
    "fig = plt.figure()\n",
    "plt.suptitle(\"Travelling sine wave\")\n",
    "ax = fig.add_subplot(111,projection='3d')\n",
    "ax.set_xlabel(\"spatial distance\")\n",
    "ax.set_ylabel('temporal distance')\n",
    "ax.set_zlabel(\"amplitude\")\n",
    "l1 = ax.plot(x[0], t[0], signal(x[0], t[0], velocities), label='right moving')\n",
    "l2 = ax.plot(x[0], t[0], signal(x[0], t[0], velocities[::-1]), label='left moving')\n",
    "l3 = ax.plot(x[0], t[0], signal(x[0], t[0], 1), label='flat line')\n",
    "\n",
    "lines = [l1, l2, l3]\n",
    "\n",
    "def update(i, fig, ax):\n",
    "    ax.view_init(elev=20., azim=i)\n",
    "    \n",
    "    \n",
    "    ax.legend()\n",
    "    return fig, ax\n",
    " \n",
    "anim = FuncAnimation(fig, update, frames=np.arange(0, 360, 2), repeat=True, fargs=(fig, ax))\n",
    "anim.save('rgb_cube.gif', dpi=80, writer='imagemagick', fps=24)\n",
    "plt.show();"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
