{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discrete Sine + Delta correlation\n",
    "\n",
    "Bandpass a delta peak, and try to correlate it with an appropriate sine wave.\n",
    "\\\n",
    "Show that we can lift the degeneracy in the correlations (within a window)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.fft as ft\n",
    "import scipy.signal as signal\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": [],
   "source": [
    "def plot_spectrum( ax, spectrum, freqs, plot_complex=False, plot_power=False, plot_amplitude=None):\n",
    "    \"\"\" Plot a signal's spectrum on an Axis object\"\"\"\n",
    "    plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)\n",
    "    alpha = 1\n",
    "    \n",
    "    ax.set_title(\"Spectrum\")\n",
    "    ax.set_xlabel(\"f (Hz)\")\n",
    "    ylabel = \"\"\n",
    "    if plot_amplitude or plot_complex:\n",
    "        ylabel = \"Amplitude\"\n",
    "    if plot_power:\n",
    "        if ylabel:\n",
    "            ylabel += \"|\"\n",
    "        ylabel += \"Power\"\n",
    "    ax.set_ylabel(ylabel)\n",
    "\n",
    "    if plot_complex:\n",
    "        alpha = 0.5\n",
    "        ax.plot(freqs, np.real(spectrum), '.-', label='Real', alpha=alpha)\n",
    "        ax.plot(freqs, np.imag(spectrum), '.-', label='Imag', alpha=alpha)\n",
    "\n",
    "    if plot_power:\n",
    "        ax.plot(freqs, np.abs(spectrum)**2, '.-', label='Power', alpha=alpha)\n",
    "    \n",
    "    if plot_amplitude:\n",
    "        ax.plot(freqs, np.abs(spectrum), '.-', label='Abs', alpha=alpha)\n",
    "\n",
    "    ax.legend()\n",
    "\n",
    "    return ax\n",
    "\n",
    "def fft_bandpass(signal, band, samplerate):\n",
    "    \"\"\"\n",
    "    Simple bandpassing function employing a FFT.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    signal : arraylike\n",
    "    band : tuple(low, high)\n",
    "        Frequencies for bandpassing\n",
    "    samplerate : float\n",
    "    \"\"\"\n",
    "    signal = np.asarray(signal)\n",
    "\n",
    "    fft = ft.rfft(signal)\n",
    "    freqs = ft.rfftfreq(signal.size, 1/samplerate)\n",
    "    fft[(freqs < band[0]) | (freqs > band[1])] = 0\n",
    "    \n",
    "    return ft.irfft(fft, signal.size), (fft, freqs)\n",
    "\n",
    "def deltapeak(timelength=1e3, samplerate=1, offset=None, peaklength=1):\n",
    "    N_samples = int(timelength * samplerate)\n",
    "    if offset is None:\n",
    "        offset = (np.random.random()*N_samples).astype(int) % N_samples\n",
    "    \n",
    "    position = (offset + np.arange(0, peaklength)).astype(int) % N_samples\n",
    "    \n",
    "    signal = np.zeros(N_samples)\n",
    "    signal[position] = 1\n",
    "    \n",
    "    return signal, position\n",
    "\n",
    "def sin_delay(f, t, t_delay=0, phase=0):\n",
    "    return np.cos(2*np.pi*f * (t - t_delay) + phase)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate a delta peak and bandpass it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size: 25; Offset: 0; PeakLength: 1\n"
     ]
    }
   ],
   "source": [
    "# Generate a delta peak and bandpass it\n",
    "samplerate = 500 # MHz\n",
    "band = (30, 80) # MHz\n",
    "\n",
    "us = 1e3 # ns\n",
    "ns = 1/us # us\n",
    "\n",
    "## Delta peak properties\n",
    "imp_timelength = 50 * ns# ns\n",
    "imp_peaklength = 1\n",
    "imp_time = np.arange(0, imp_timelength, 1/samplerate)\n",
    "imp_offset = [1400]\n",
    "\n",
    "orig_imp, imp_offset = deltapeak(imp_timelength, samplerate, offset=imp_offset, peaklength=imp_peaklength)\n",
    "\n",
    "## Bandpass it\n",
    "imp, (fft, fft_freqs) = fft_bandpass(orig_imp, band, samplerate)\n",
    "imp /= np.max(imp)\n",
    "\n",
    "#\n",
    "print(\"Size: {}; Offset: {}; PeakLength: {}\".format(orig_imp.size, imp_offset[0], len(imp_offset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-4-5b3a8514ce9c>:22: UserWarning: Matplotlib is currently using module://ipykernel.pylab.backend_inline, which is a non-GUI backend, so cannot show the figure.\n",
      "  fig.show()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37500 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot timedomain and frequency domain of the bandpassed signal\n",
    "fig, (ax2, ax1) = plt.subplots(1,2, figsize=(16,4),sharey=True)\n",
    "## Spectrum\n",
    "plot_spectrum(ax1, fft, fft_freqs)\n",
    "ax1.set_xlabel(\"Frequency [MHz]\")\n",
    "ax1.grid(True)\n",
    "ax1.margins(0.1, 0.1)\n",
    "ax1.get_legend().remove()\n",
    "ax1.set_xlim(0, 200)\n",
    "\n",
    "## Signal\n",
    "ax2.set_title(\"Signal\")\n",
    "ax2.margins(0.1, 0.1)\n",
    "ax2.set_xlabel('Time [ns]')\n",
    "ax2.set_ylabel('Amplitude')\n",
    "ax2.grid(True)\n",
    "ax2.plot(imp_time/ns, imp, label='Bandpassed')\n",
    "ax2.plot(imp_time/ns, orig_imp, label='Original', zorder=-1, color='g')\n",
    "ax2.axvline(imp_time[imp_offset[0]]/ns, color='r', ls='--')\n",
    "ax2.legend()\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "### correlation\n",
    "def correlation_and_lag(sig1, sig2, mode=\"same\", normalise=False):\n",
    "    corr = signal.correlate(sig1, sig2, mode=mode)\n",
    "    if normalise:\n",
    "        corr /= np.max(corr)\n",
    "\n",
    "    lags = signal.correlation_lags(sig1.size, sig2.size, mode=mode)\n",
    "    \n",
    "    return corr, lags\n",
    "\n",
    "def find_best_lag(sig1, sig2, fix_one_short=False, fix_positive=False, **corr_kwargs):\n",
    "    corr, lags = correlation_and_lag(sig1, sig2, **corr_kwargs)\n",
    "    \n",
    "    lag = lags[np.argmax(corr)]\n",
    "    \n",
    "    if fix_one_short:\n",
    "        # for some reason it is always one short\n",
    "        if lag > 0:\n",
    "            lag += 1\n",
    "        elif lag < 0:\n",
    "            lag -= 1\n",
    "\n",
    "    # turn negative lag into positive\n",
    "    if fix_positive and lag < 0:\n",
    "        print(\"negative lag\")\n",
    "        lag += len(sig2)\n",
    "    \n",
    "    return lag, (corr, lags)\n",
    "\n",
    "\n",
    "def find_beacon_phase_delay(samplerate, f_beacon, reference_beacon, delayed_beacon, **lag_kwargs):\n",
    "    \"\"\"\n",
    "    Return phase delay of `beacon` with respect to `reference_beacon`.\n",
    "    Note that the returned value can be off by a multiple of $2\\pi$.\n",
    "    \n",
    "    Parameters\n",
    "    ==========\n",
    "    samplerate : float\n",
    "        Samplerate of both reference_beacon and delayed_beacon\n",
    "    f_beacon : float\n",
    "        Frequency of the beacons\n",
    "    reference_beacon : ndarray\n",
    "        The beacon to use as a reference\n",
    "    beacon : ndarray\n",
    "        The beacon to find the delay for\n",
    "    \"\"\"\n",
    "    \n",
    "    calc_lag, _ = find_best_lag(reference_beacon, delayed_beacon, **lag_kwargs)\n",
    "    \n",
    "    return 2*np.pi* f_beacon * calc_lag /samplerate \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Correlate a sinusoid with the bandpassed delta peak"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Known offset [ns]: 6.0\n",
      "Known phase [ns]: 1.884955592153876\n",
      "Expected lag: 3.0\n",
      "Lag time multiplier [ns]: 2.0\n",
      "Calc Lag: 1\n",
      "Calc Lag time [ns]: 2.0\n",
      "Phase lag: 0.6283185307179586\n"
     ]
    }
   ],
   "source": [
    "###########################################\n",
    "# Correlate a sinusoid\n",
    "\n",
    "## define sinusoid properties\n",
    "\n",
    "# in-band or out-of-band\n",
    "if True:\n",
    "    f_sine = 50 # MHz\n",
    "    corr_sig_timelength = 1/f_sine # us\n",
    "else:\n",
    "    f_sine = 20 # MHz\n",
    "    corr_sig_timelength = 1/f_sine # us\n",
    "    \n",
    "## define general correlation signal properties\n",
    "corr_sig_samplerate = samplerate # MHz\n",
    "corr_sig_samplelength = corr_sig_timelength * corr_sig_samplerate\n",
    "corr_time = np.arange(0, corr_sig_timelength, 1/corr_sig_samplerate)\n",
    "\n",
    "corr_sig = sin_delay(f_sine, corr_time, t_delay=0, phase=-1/2*np.pi)\n",
    "\n",
    "## define the signals\n",
    "if not True:\n",
    "    sig1 = imp\n",
    "    sig1_time = imp_time\n",
    "    sig1_offset = imp_time[imp_offset[0]]\n",
    "else:\n",
    "    sig1_offset = 6*ns # us\n",
    "    sig1 = sin_delay(f_sine, corr_time, t_delay=sig1_offset)\n",
    "    sig1_time = corr_time\n",
    "\n",
    "\n",
    "## correlate it and find the best correlation\n",
    "calc_lag, (corr, lags) = find_best_lag(sig1, corr_sig, mode='full', fix_one_short=False)\n",
    "\n",
    "calc_lag += 1\n",
    "\n",
    "#####\n",
    "lag_times = lags / corr_sig_samplerate # us\n",
    "calc_lag_time = (calc_lag) / corr_sig_samplerate # us\n",
    "\n",
    "\n",
    "\n",
    "#(calc_lag_time - 1/corr_sig_samplerate)\n",
    "#((calc_lag_time*corr_sig_samplerate - 1) - (calc_lag_time - 1)) /samplerate\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#calc_lag_time = ((len(corr)-1)/2 - corr.argmax())/corr_sig_samplerate\n",
    "\n",
    "calc_phase_lag = 2*np.pi*f_sine*calc_lag_time\n",
    "\n",
    "print(\"Known offset [ns]:\", sig1_offset/ns)\n",
    "print(\"Known phase [ns]:\", 2*np.pi*f_sine*sig1_offset)\n",
    "print(\"Expected lag:\", sig1_offset*corr_sig_samplerate)\n",
    "print(\"Lag time multiplier [ns]:\", 1/corr_sig_samplerate/ns)\n",
    "print(\"Calc Lag:\", calc_lag)\n",
    "print(\"Calc Lag time [ns]:\", calc_lag_time/ns)\n",
    "print(\"Phase lag:\", calc_phase_lag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-7-42b7d03bc798>:26: UserWarning: Matplotlib is currently using module://ipykernel.pylab.backend_inline, which is a non-GUI backend, so cannot show the figure.\n",
      "  fig.show();\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33568 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#####\n",
    "# Plot signal\n",
    "#####\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlabel(\"Time [ns]\")\n",
    "ax.set_ylabel(\"Amplitude\")\n",
    "\n",
    "\n",
    "ax.plot(sig1_time/ns, sig1+2, label='Signal')\n",
    "ax.axvline(sig1_offset/ns, color='r', ls='--', label='offset')\n",
    "\n",
    "\n",
    "ax.axvline(calc_lag_time/ns, color='b', ls=(2, (10, 10)), label='calc offset')\n",
    "ax.plot(corr_time/ns, corr_sig+2, label=\"Uncorr\")\n",
    "ax.plot(sig1_time/ns, sin_delay(f_sine, sig1_time, t_delay=calc_lag_time), label=\"Corr time\")\n",
    "ax.plot(sig1_time/ns, sin_delay(f_sine, sig1_time, phase=-1*calc_phase_lag)-0.2, label=\"Corr phase\")\n",
    "\n",
    "# debugging\n",
    "if True:\n",
    "    ax.plot(corr_time[-1]/ns + corr_time/ns, sin_delay(f_sine, corr_time, t_delay=sig1_offset)+2, label=\"Cheat time\")\n",
    "    ax.plot(corr_time[-1]/ns + corr_time/ns, sin_delay(f_sine, corr_time, phase=-2*np.pi*sig1_offset*f_sine)+2, label=\"Cheat phase\")\n",
    "    \n",
    "    \n",
    "ax.legend()\n",
    "fig.show();\n",
    "print((sig1_offset - calc_lag_time)*samplerate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-8-b80dbc5f3c93>:20: UserWarning: Matplotlib is currently using module://ipykernel.pylab.backend_inline, which is a non-GUI backend, so cannot show the figure.\n",
      "  fig.show();\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23448 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#####\n",
    "# Plot Correlation\n",
    "#####\n",
    "max_corr_id = np.argmax(corr)\n",
    "min_corr_id = np.argmin(corr)\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_xlabel(\"Lag [ns]\")\n",
    "ax.set_ylabel(\"Correlation\")\n",
    "ax.plot(lag_times/ns, corr)\n",
    "\n",
    "ax.plot()\n",
    "\n",
    "ax.axvline(lag_times[max_corr_id]/ns, ls='--',color='g', label=\"max_corr\")\n",
    "ax.axhline(corr[max_corr_id], ls='--', color='g')\n",
    "ax.axvline(lag_times[min_corr_id]/ns, ls='--', color='r', label='min_corr')\n",
    "ax.axhline(corr[min_corr_id], ls='--', color='r')\n",
    "\n",
    "ax.axvline(calc_lag_time/ns, ls=(0, (5, 5)), color='purple', label='calculated')\n",
    "ax.legend()\n",
    "fig.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
}
