{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Perform a grid search to find the timing error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "# Append parent directory to import path so lib can be found\n",
    "sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))\n",
    "\n",
    "import numpy as np\n",
    "rng = np.random.default_rng(12345)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import signal\n",
    "import copy\n",
    "\n",
    "from lib.signals import DigitisedSignal\n",
    "import lib.util as util\n",
    "\n",
    "\n",
    "km = 1e3\n",
    "ns = 1e9\n",
    "\n",
    "# monkey patch correlation_lags into signal if it does not exist\n",
    "if not hasattr(signal, 'correlation_lags'):\n",
    "    def correlation_lags(in1_len, in2_len, mode='full'):\n",
    "        r\"\"\"\n",
    "        Calculates the lag / displacement indices array for 1D cross-correlation.\n",
    "        Parameters\n",
    "        ----------\n",
    "        in1_size : int\n",
    "            First input size.\n",
    "        in2_size : int\n",
    "            Second input size.\n",
    "        mode : str {'full', 'valid', 'same'}, optional\n",
    "            A string indicating the size of the output.\n",
    "            See the documentation `correlate` for more information.\n",
    "        See Also\n",
    "        --------\n",
    "        correlate : Compute the N-dimensional cross-correlation.\n",
    "        Returns\n",
    "        -------\n",
    "        lags : array\n",
    "            Returns an array containing cross-correlation lag/displacement indices.\n",
    "            Indices can be indexed with the np.argmax of the correlation to return\n",
    "            the lag/displacement.\n",
    "        Notes\n",
    "        -----\n",
    "        Cross-correlation for continuous functions :math:`f` and :math:`g` is\n",
    "        defined as:\n",
    "        .. math::\n",
    "            \\left ( f\\star g \\right )\\left ( \\tau \\right )\n",
    "            \\triangleq \\int_{t_0}^{t_0 +T}\n",
    "            \\overline{f\\left ( t \\right )}g\\left ( t+\\tau \\right )dt\n",
    "        Where :math:`\\tau` is defined as the displacement, also known as the lag.\n",
    "        Cross correlation for discrete functions :math:`f` and :math:`g` is\n",
    "        defined as:\n",
    "        .. math::\n",
    "            \\left ( f\\star g \\right )\\left [ n \\right ]\n",
    "            \\triangleq \\sum_{-\\infty}^{\\infty}\n",
    "            \\overline{f\\left [ m \\right ]}g\\left [ m+n \\right ]\n",
    "        Where :math:`n` is the lag.\n",
    "        Examples\n",
    "        --------\n",
    "        Cross-correlation of a signal with its time-delayed self.\n",
    "        >>> from scipy import signal\n",
    "        >>> from numpy.random import default_rng\n",
    "        >>> rng = default_rng()\n",
    "        >>> x = rng.standard_normal(1000)\n",
    "        >>> y = np.concatenate([rng.standard_normal(100), x])\n",
    "        >>> correlation = signal.correlate(x, y, mode=\"full\")\n",
    "        >>> lags = signal.correlation_lags(x.size, y.size, mode=\"full\")\n",
    "        >>> lag = lags[np.argmax(correlation)]\n",
    "        \"\"\"\n",
    "\n",
    "        # calculate lag ranges in different modes of operation\n",
    "        if mode == \"full\":\n",
    "            # the output is the full discrete linear convolution\n",
    "            # of the inputs. (Default)\n",
    "            lags = np.arange(-in2_len + 1, in1_len)\n",
    "        elif mode == \"same\":\n",
    "            # the output is the same size as `in1`, centered\n",
    "            # with respect to the 'full' output.\n",
    "            # calculate the full output\n",
    "            lags = np.arange(-in2_len + 1, in1_len)\n",
    "            # determine the midpoint in the full output\n",
    "            mid = lags.size // 2\n",
    "            # determine lag_bound to be used with respect\n",
    "            # to the midpoint\n",
    "            lag_bound = in1_len // 2\n",
    "            # calculate lag ranges for even and odd scenarios\n",
    "            if in1_len % 2 == 0:\n",
    "                lags = lags[(mid-lag_bound):(mid+lag_bound)]\n",
    "            else:\n",
    "                lags = lags[(mid-lag_bound):(mid+lag_bound)+1]\n",
    "        elif mode == \"valid\":\n",
    "            # the output consists only of those elements that do not\n",
    "            # rely on the zero-padding. In 'valid' mode, either `in1` or `in2`\n",
    "            # must be at least as large as the other in every dimension.\n",
    "\n",
    "            # the lag_bound will be either negative or positive\n",
    "            # this let's us infer how to present the lag range\n",
    "            lag_bound = in1_len - in2_len\n",
    "            if lag_bound >= 0:\n",
    "                lags = np.arange(lag_bound + 1)\n",
    "            else:\n",
    "                lags = np.arange(lag_bound, 1)\n",
    "        return lags\n",
    "\n",
    "    signal.correlation_lags = correlation_lags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51304 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Setup signal\n",
    "########################\n",
    "\n",
    "signal_t_length = 200e-9\n",
    "signal_sample_rate = 5e9 # Hz\n",
    "signal_time = util.sampled_time(\n",
    "    sample_rate=signal_sample_rate,\n",
    "    start=0,\n",
    "    end=signal_t_length\n",
    ")\n",
    "double_signal_time = np.concatenate((signal_time, signal_time[-1] + signal_time)) \\\n",
    "                     - (signal_time[-1] - signal_time[0])/2\n",
    "\n",
    "fc = 50e6# MHz\n",
    "\n",
    "if True:\n",
    "\n",
    "    shift = -25.8e-9\n",
    "    signal_pulse1_center = 125e-9 # 100 ns\n",
    "    \n",
    "    if not True:\n",
    "        signal_1 = DigitisedSignal(signal.gausspulse(signal_time - signal_pulse1_center, fc=fc),\n",
    "                                   signal_sample_rate)\n",
    "        signal_2 = DigitisedSignal(signal.gausspulse(signal_time - signal_pulse1_center, fc=fc), \n",
    "                                   signal_sample_rate)\n",
    "        signal_2.t_0 = shift\n",
    "    else:\n",
    "        signal_1 = DigitisedSignal(np.sin(2*np.pi*(signal_time)*fc/5.4) + np.sin(2*np.pi*(signal_time)*fc),\n",
    "                                   signal_sample_rate)\n",
    "\n",
    "        signal_2 = copy.copy(signal_1)\n",
    "        signal_2.t_0 = shift\n",
    "   \n",
    "    signal_1.periodic = False\n",
    "    signal_2.periodic = False\n",
    "    del signal_pulse1_center\n",
    "\n",
    "\n",
    "del fc, signal_t_length, signal_sample_rate\n",
    "\n",
    "\n",
    "# Pre calculated values\n",
    "sig_1 = signal_1(double_signal_time)\n",
    "sig_2 = signal_2(double_signal_time)\n",
    "\n",
    "\n",
    "if True:\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(double_signal_time * ns, sig_1, label='Regular')\n",
    "    ax.plot((double_signal_time) * ns, sig_2, label=\"Shifted\")\n",
    "    ax.plot((double_signal_time) * ns, signal_2(double_signal_time + shift), 'r--', label=\"Cheat\")\n",
    "    ax.set_xlabel(\"Time (ns)\")\n",
    "    ax.set_ylabel(\"Amplitude\")\n",
    "    ax.set_title(\"Lag: {:.4g} ns\".format(shift*ns))\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52384 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Using scipy\n",
    "\n",
    "if True:\n",
    "    scipy_corr = signal.correlate(sig_1, sig_2, mode='same')\n",
    "    scipy_id = scipy_corr.argmax()\n",
    "    calc_shift =  ((len(scipy_corr)-1)/2 -scipy_id)/signal_1.sample_rate\n",
    "\n",
    "    if not True:\n",
    "        signal_2.periodic = True\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(double_signal_time * ns, sig_1, label='Regular')\n",
    "    ax.plot((double_signal_time) * ns, sig_2, label=\"Shifted\")\n",
    "    ax.plot((double_signal_time) * ns, signal_2(double_signal_time + calc_shift), 'r--', label='Recovered')\n",
    "    ax.set_xlabel(\"Time (ns)\")\n",
    "    ax.set_ylabel(\"Correlation\")\n",
    "    ax.set_title(\"Calculated Lag: {:.4g} ns\".format(calc_shift*ns))\n",
    "    ax.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## new code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lag_gridsearch(grid, sample_rate, signal_data, reference):\n",
    "    \"\"\"\n",
    "    Return the best time shift found when doing a grid search.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    lag_grid - ndarray\n",
    "        The array specifying the grid that is to be searched.\n",
    "    sample_rate - float\n",
    "        Sample rate of signal_data to transform index to time.\n",
    "    signal_data - ndarray\n",
    "        The real signal to find the time shift for.\n",
    "    reference - ndarray\n",
    "        Real signal to use as reference to obtain lag.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    lag : ndarray\n",
    "        The best time shift obtained\n",
    "    err : tuple\n",
    "        Difference to the previous and next time shift from lag, resp.\n",
    "    \"\"\"\n",
    "\n",
    "    assert signal_data.shape == reference.shape\n",
    "    \n",
    "    corrs = grid_correlate(grid, signal_data, reference)\n",
    "    \n",
    "    idx = np.argmax(corrs)\n",
    "    \n",
    "    lag = grid[idx]/sample_rate\n",
    "    \n",
    "    err_min =  (grid[idx-1]-grid[idx])/(2*sample_rate)\n",
    "    err_plus = (grid[idx+1]-grid[idx])/(2*sample_rate)\n",
    "\n",
    "    return lag, (err_min, err_plus)\n",
    "    \n",
    "\n",
    "def grid_correlate(grid, x, reference):\n",
    "    \"\"\"\n",
    "    Determine correlation between x and reference using grid as \n",
    "    the lags to be used for the correlation.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    grid - ndarray\n",
    "        The array specifying the grid that is to be searched.\n",
    "    x - ndarray\n",
    "        The real signal to find the time shift for.\n",
    "    reference - ndarray\n",
    "        Real signal to use as reference to obtain lag.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    corrs - ndarray\n",
    "        The correlations along grid.\n",
    "    \"\"\"\n",
    "    grid = np.asarray(grid)\n",
    "    x = np.asarray(x)\n",
    "    reference = np.asarray(reference)\n",
    "\n",
    "    assert x.shape == reference.shape\n",
    "    \n",
    "    ref_conj = np.conjugate(reference)\n",
    "    \n",
    "    corrs = np.array([np.dot(np.roll(ref_conj, lag), x) for lag in grid], dtype=np.float64)\n",
    "    \n",
    "    return corrs\n",
    "\n",
    "def correlation_grid(grid_size=None, in1_len=None, in2_len = None, end = None, start=None, mode='full'):\n",
    "    \"\"\"\n",
    "    Abuse correlation_lags to determine the endpoints of the grid.\n",
    "    \"\"\"\n",
    "    \n",
    "    if in1_len is not None or in2_len is not None:\n",
    "        if in2_len is None:\n",
    "            in2_len = in1_len\n",
    "        elif in1_len is None:\n",
    "            in1_len = in2_len\n",
    "\n",
    "        lags = signal.correlation_lags(in1_len, in2_len, mode=mode)\n",
    "\n",
    "        max_lag = max(lags)\n",
    "        min_lag = min(lags)\n",
    "    else:\n",
    "        max_lag = np.inf\n",
    "        min_lag = -np.inf\n",
    "\n",
    "    if end is None:\n",
    "        end = max_lag\n",
    "    elif end > max_lag:\n",
    "        raise ValueError(\"Grid end is too high\")\n",
    "\n",
    "    if start is None:\n",
    "        start = min_lag\n",
    "    elif start < min_lag:\n",
    "        raise ValueError(\"Grid start is too low\")\n",
    "\n",
    "    return np.linspace(start, end, grid_size, dtype=int, endpoint=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debug and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39308 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ids 87 871\n",
      "ts -2.62e-08 -2.58e-08 =: -2.58e-08\n",
      "avail ts -2e-07 - 1.978e-07 : 1.989e-09\n",
      "grid size 200 2000 2000\n"
     ]
    }
   ],
   "source": [
    "#################\n",
    "grid_size = 200\n",
    "my_mode = 'same'\n",
    "grid = correlation_grid(grid_size, in1_len=len(sig_1), in2_len=len(sig_2), mode=my_mode)\n",
    "t_grid = grid/signal_1.sample_rate\n",
    "\n",
    "scipy_mode = my_mode\n",
    "scipy_corr = signal.correlate(sig_2, sig_1, mode=scipy_mode)\n",
    "scipy_lags = signal.correlation_lags(len(sig_2), len(sig_1), mode=scipy_mode)\n",
    "corrs = grid_correlate(grid, sig_2, sig_1)\n",
    "   \n",
    "\n",
    "scipy_id = np.argmax(scipy_corr)\n",
    "my_id = np.argmax(corrs)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_title(\"Correlations\")\n",
    "\n",
    "if not True:\n",
    "    x_grid = grid\n",
    "    scipy_grid = scipy_lags\n",
    "    ax.set_xlabel(\"Index Lag\")\n",
    "else:\n",
    "    x_grid = t_grid * 1e9\n",
    "    scipy_grid = scipy_lags / signal_1.sample_rate * 1e9\n",
    "    ax.set_xlabel(\"Time [ns]\")\n",
    "    \n",
    "ax.plot(x_grid, corrs, label='grid correlation')\n",
    "ax.plot(scipy_grid, scipy_corr, label='scipy')\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "plt.show();\n",
    "\n",
    "\n",
    "if scipy_mode == 'same':\n",
    "    scipy_corr_grid = scipy_corr[grid]\n",
    "else:\n",
    "    scipy_corr_grid = scipy_corr[grid]\n",
    "    \n",
    "\n",
    "if len(grid) <= 10:\n",
    "    print(\"G\", grid)\n",
    "    print(\"TG\", grid/signal_1.sample_rate)\n",
    "    print(\"C\", corrs)\n",
    "    print(\"SC\", scipy_corr_grid)\n",
    "    \n",
    "    \n",
    "print(\"ids\", my_id, scipy_id)\n",
    "print(\"ts\", t_grid[my_id], scipy_lags[np.argmax(scipy_corr)]/signal_1.sample_rate, '=:', signal_2.t_0)\n",
    "print(\"avail ts\", t_grid[0], '-', t_grid[-1] ,':', (t_grid[-1] - t_grid[0])/len(t_grid))\n",
    "print(\"grid size\", len(t_grid), len(sig_2), len(sig_1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Showcase functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Known time offset: -2.58e-08\n",
      "Calculated offset: -2.62e-08\n",
      "Uncertainty on offset: (-1e-09, 1e-09)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 64756 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "grid_size = 200\n",
    "\n",
    "grid = correlation_grid(grid_size, in1_len=len(sig_1), in2_len=len(sig_2), mode='same')\n",
    "calc_shift, err = lag_gridsearch(grid, signal_1.sample_rate, sig_2, sig_1)\n",
    "\n",
    "print(\"Known time offset:\", signal_2.t_0)\n",
    "print(\"Calculated offset:\", calc_shift)\n",
    "print(\"Uncertainty on offset:\", err)\n",
    "\n",
    "if True:\n",
    "    signal_2.periodic = False\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16,4))\n",
    "ax.plot(double_signal_time   * ns, sig_1, alpha=0.5, label='Regular')\n",
    "ax.plot(double_signal_time   * ns, sig_2, alpha=0.5, label='Shifted')\n",
    "ax.plot((double_signal_time) * ns, signal_2(double_signal_time + calc_shift), ls=(0, (10,)), label='Recovered')\n",
    "\n",
    "\n",
    "ax.axvline(calc_shift * ns)\n",
    "for i, _ in enumerate(err):\n",
    "    ax.axvline((calc_shift + err[i]) * ns, ls=(5*i, (4, 4)))\n",
    "\n",
    "ax.set_xlabel(\"Time (ns)\")\n",
    "ax.set_ylabel(\"Amplitude\")\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.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
