{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Translation of a discrete 1D signal\n",
    "\n",
    "\n",
    "#### wave equation\n",
    "\n",
    "$$\n",
    "\\frac{\\partial^2 u}{\\partial t^2}\n",
    " = \n",
    " v^2 \\frac{\\partial^2 u}{\\partial x^2}\n",
    "$$\n",
    "\n",
    "\n",
    "##### simple sine solution\n",
    "\n",
    "$$\n",
    "u(x, t) = \\sin\\left( \\frac{2 \\pi x}{\\lambda} -2\\pi f t\\right)\n",
    "= \\sin\\left( 2 \\pi f \\, ( \\frac{x}{v(f)} -t) \\right)\n",
    "$$\n",
    "\n",
    "$$\\lambda = \\dfrac{v}{f}$$\n",
    "##### plane wave per mode\n",
    "\n",
    "\n",
    "\n",
    "$$ \n",
    "u_f(x,t) \n",
    "= \n",
    " A \\exp(- 2\\pi i f \\left[ \\frac{x}{v(f)} - t \\right] )\n",
    " +\n",
    " B \\exp(- 2\\pi i f \\left[ \\frac{x}{v(f)} + t \\right] )\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import scipy.fft as ft\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import matplotlib.ticker as tck\n",
    "\n",
    "rng = np.random.default_rng()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TravelSignal:\n",
    "    \"\"\"\n",
    "    Model an arbitrary digitised signal that can be translated to another position and time.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, signal, sample_rate, t_0 = 0, x_0 = 0, periodic=True, interp1d_kw = None):\n",
    "        \"\"\"\n",
    "        Initialise by saving the raw signal\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        signal : arraylike\n",
    "            The raw signal to wrap.\n",
    "        sample_rate : float\n",
    "            Sample rate of the raw signal.\n",
    "        t_0 : float, optional\n",
    "            Time that this signal is sent out.\n",
    "        x_0 : float, optional\n",
    "            Location that this signal is sent out from.\n",
    "        periodic : bool, optional\n",
    "            Translated signal is 0 if it is not periodic\n",
    "            and the time/distance is outside the samples.\n",
    "        interp1d_kw : bool or dict, optional\n",
    "            Use scipy.interpolate's interp1d_kw for interpolation.\n",
    "            Set to True, or a dictionary to enable.\n",
    "            Dictionary will be entered in as **kwargs.\n",
    "        \"\"\"\n",
    "\n",
    "        self.raw = signal\n",
    "        self.periodic = periodic\n",
    "\n",
    "        self.sample_rate = sample_rate # Hz\n",
    "        self.sample_length = len(self.raw)\n",
    "        self.time_length = self.sample_length*sample_rate # s\n",
    "        \n",
    "        self.x_0 = x_0\n",
    "        self.t_0 = t_0\n",
    "\n",
    "        # choose interpolation method\n",
    "        if not interp1d_kw:\n",
    "            self.interp_f = None\n",
    "\n",
    "        else:\n",
    "            # offload interpolation to scipy.interpolate\n",
    "            import scipy.interpolate as interp\n",
    "\n",
    "            interp1d_kw_defaults = {\n",
    "                \"copy\": False,\n",
    "                \"kind\": 'linear',\n",
    "                \"assume_sorted\": True,\n",
    "                \"bounds_error\": True\n",
    "            }\n",
    "\n",
    "            if self.periodic:\n",
    "                interp1d_kw_defaults['bounds_error'] = False\n",
    "                interp1d_kw_defaults['fill_value'] = (self.raw[-1], self.raw[0])\n",
    "                \n",
    "            # merge kwargs\n",
    "            if interp1d_kw is not True:\n",
    "                interp1d_kw = { **interp1d_kw_defaults, **interp1d_kw }\n",
    "\n",
    "            self.interp_f = interp.interp1d(\n",
    "                                np.arange(0, self.sample_length),\n",
    "                                self.raw,\n",
    "                                **interp1d_kw\n",
    "                            )\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.sample_length\n",
    "        \n",
    "    def __call__(self, t_f = None, x_f = None, **kwargs):\n",
    "        \"\"\"\n",
    "        Allow this class to be used as a function.\n",
    "        \"\"\"\n",
    "        return self._translate(t_f, x_f, **kwargs)[0]\n",
    "    \n",
    "    def _translate(self, t_f = None, x_f = None, t_0 = None, x_0 = None, velocity = None):\n",
    "        \"\"\"\n",
    "        Translate the signal from (t_0, x_0) to (t_f, x_f) with optional velocity.\n",
    "        \n",
    "        Returns the signal at (t_f, x_f)\n",
    "        \"\"\"\n",
    "        \n",
    "        if t_0 is None:\n",
    "            t_0 = self.t_0\n",
    "        \n",
    "        if velocity is None:\n",
    "            velocity = 1\n",
    "\n",
    "\n",
    "        ## spatial offset\n",
    "        if x_f is None:\n",
    "            spatial_time_offset = 0\n",
    "        else:\n",
    "            x_f = np.asarray(x_f)\n",
    "            if x_0 is None:\n",
    "                x_0 = self.x_0\n",
    "\n",
    "            spatial_time_offset = np.sum(np.sqrt( (x_f - x_0)**2 )/velocity)\n",
    "\n",
    "        ## temporal offset\n",
    "        if t_f is None:\n",
    "            temporal_time_offset = 0\n",
    "        else:\n",
    "            t_f = np.asarray(t_f)\n",
    "                \n",
    "            if t_0 is None:\n",
    "                t_0 = self.t_0\n",
    "        \n",
    "            temporal_time_offset = t_f - t_0\n",
    "\n",
    "        # total offset\n",
    "        total_time_offset = spatial_time_offset + temporal_time_offset\n",
    "        n_offset = (total_time_offset * sample_rate )\n",
    "\n",
    "        # periodic signal\n",
    "        if self.periodic:\n",
    "            n_offset = n_offset % self.sample_length\n",
    "\n",
    "        # non-periodic and outside the bounds\n",
    "        else:\n",
    "            mask_idx = np.nonzero( (0 > n_offset) | (n_offset >= self.sample_length) )\n",
    "            n_offset[mask_idx] = 0\n",
    "\n",
    "        # offload to scipy interpolation\n",
    "        if self.interp_f:\n",
    "            amplitude = self.interp_f(n_offset)\n",
    "        \n",
    "        # self written linear interpolation\n",
    "        else:\n",
    "            n_offset_eps, n_offset_int = np.modf(n_offset)\n",
    "            n_offset_int = n_offset.astype(int)\n",
    "\n",
    "            if True:\n",
    "                amplitude = (1-n_offset_eps) * self.raw[n_offset_int] \\\n",
    "                            + n_offset_eps * self.raw[(n_offset_int + 1) % self.sample_length]\n",
    "\n",
    "            # use nearest value instead of interpolation\n",
    "            else:\n",
    "                amplitude = self.raw[n_offset_int]\n",
    "\n",
    "        if not self.periodic:\n",
    "            amplitude[mask_idx] = 0\n",
    "        \n",
    "        return amplitude, total_time_offset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolation: Sine wave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55092 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Sine wave\n",
    "sample_rate = 1e3 # Hz\n",
    "f_sine = 10\n",
    "time   = np.arange(0, 1, 1/(sample_rate)) #s\n",
    "time2   = np.arange(0, 1, 1/(sample_rate*2)) #s\n",
    "phase_offset = 0\n",
    "\n",
    "signal_func = lambda t: np.cos(2*np.pi*f_sine*t + phase_offset)\n",
    "\n",
    "mysignal = TravelSignal(signal_func(time), sample_rate)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16,4))\n",
    "ax.plot(time,  signal_func(time),    '.-',  label=\"Raw signal\")\n",
    "ax.plot(time,  mysignal(time)+0.5,   '.-',  label=\"MySignal +0.5\")\n",
    "ax.plot(time2, mysignal(time2)+0.25, '.--', label=\"MySignal(interp) +0.25\")\n",
    "ax.set_ylabel(\"Amplitude\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.set_xlim(0, 1/f_sine)\n",
    "ax.legend()\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extrapolation: Sine wave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 98140 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Sine wave\n",
    "sample_rate = 1e3 # Hz\n",
    "f_sine = 10\n",
    "time   = np.arange(0, 1, 1/(sample_rate)) #s\n",
    "time2   = np.arange(3/4, 2, 1/(sample_rate*2)) #s\n",
    "phase_offset = 0\n",
    "\n",
    "signal_func = lambda t: np.cos(2*np.pi*f_sine*t + phase_offset)\n",
    "\n",
    "mysignal = TravelSignal(signal_func(time), sample_rate)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16,4))\n",
    "ax.plot(time,  signal_func(time),    '.-',  label=\"Raw signal\")\n",
    "ax.plot(time,  mysignal(time)+0.5,   '.-',  label=\"MySignal +0.5\")\n",
    "ax.plot(time2, mysignal(time2)+0.25, '.--', label=\"MySignal(interp) +0.25\")\n",
    "ax.set_ylabel(\"Amplitude\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.legend()\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolation: Delta function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35560 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## 1d signal: Delta peak with time interval\n",
    "sample_rate = 3e1 # Hz\n",
    "interp_sample_rate = 2.3*sample_rate\n",
    "time   = np.arange(0, 1, 1/sample_rate) #s\n",
    "time2  = np.arange(0, 2, 1/(interp_sample_rate)) #s\n",
    "\n",
    "signal = np.zeros(len(time))\n",
    "if False:\n",
    "    signal[len(signal)//2] = 1\n",
    "else:\n",
    "    signal[len(signal)//2-1] = 1\n",
    "    signal[len(signal)//2] = 0.75\n",
    "    signal[len(signal)//2+1] = 1\n",
    "\n",
    "mysignal = TravelSignal(signal, sample_rate)\n",
    "time_signal,  time_offsets  = mysignal._translate(time)\n",
    "time2_signal, time2_offsets = mysignal._translate(time2)\n",
    "\n",
    "raw_max_idx = np.argmax(signal)\n",
    "my_max_idx  = np.argmax(mysignal(time))\n",
    "my2_max_idx = np.argmax(mysignal(time2))\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1,1,sharex=True, figsize=(16,4))\n",
    "ax.plot(time,  mysignal.raw,  '.-', label=\"Raw signal\")\n",
    "ax.plot(time,  time_signal,   '-',  label=\"MySignal +0.5\")\n",
    "ax.plot(time2, time2_signal,  '--', label=\"MySignal(interp) +0.25\")\n",
    "ax.plot(np.array([0, 1/sample_rate]) + time[raw_max_idx], [1.04, 1.04], label=\"Orig Samplerate\")\n",
    "ax.plot(np.array([0, 1/interp_sample_rate]) + time[raw_max_idx], [1.02, 1.02], label=\"Interp Samplerate\")\n",
    "ax.set_ylabel(\"Amplitude\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.legend()\n",
    "\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Non-Periodicity: Gaussian pulse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30216 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sample_rate = 3e2 # Hz\n",
    "interp_sample_rate = sample_rate * 1/10 # Hz\n",
    "\n",
    "t_offset = 1\n",
    "periodic = False\n",
    "\n",
    "time   = t_offset + np.arange(0, 1, 1/sample_rate) #s\n",
    "time2  = np.arange(0, 3*t_offset, 1/(interp_sample_rate)) #s\n",
    "\n",
    "def gaussian(x, mu, sigma ):\n",
    "    return 1/(sigma*np.sqrt(2*np.pi)) * np.exp(- (x-mu)**2 / sigma**2 )\n",
    "\n",
    "signal = gaussian(time, time[len(time)//2], (time[-1] - time[0])/10)\n",
    "\n",
    "\n",
    "mysignal = TravelSignal(signal, sample_rate, t_0 = t_offset, periodic=periodic)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(16,4))\n",
    "ax.set_ylabel(\"Amplitude\")\n",
    "ax.set_xlabel(\"Time\")\n",
    "\n",
    "ax.plot(time,  signal,                       label='Raw signal')\n",
    "ax.plot(time,  mysignal(time)+0.5,     '-',  label=\"MySignal +0.5\")\n",
    "ax.plot( time2, mysignal(time2)+0.25 , '--', label=\"MySignal(interp) +0.25\")\n",
    "ax.legend()\n",
    "\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
}
