{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "## Import required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "                    import math\n",
    "import sqlite3\n",
    "import struct\n",
    "import datetime\n",
    "import json\n",
    "\n",
    "import scipy.fftpack\n",
    "from scipy import signal as sig\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import patches\n",
    "from matplotlib import dates\n",
    "import numpy as np\n",
    "from scipy import signal, optimize, interpolate\n",
    "from tqdm.notebook import tnrange, tqdm\n",
    "from IPython.display import set_matplotlib_formats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_147503/2251044864.py:2: DeprecationWarning: `set_matplotlib_formats` is deprecated since IPython 7.23, directly use `matplotlib_inline.backend_inline.set_matplotlib_formats()`\n",
      "  set_matplotlib_formats('png', 'pdf')\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "set_matplotlib_formats('png', 'pdf')\n",
    "font = {'family' : 'normal',\n",
    "        'weight' : 'normal',\n",
    "        'size'   : 8}\n",
    "matplotlib.rc('font', **font)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data series information from sqlite capture file\n",
    "\n",
    "One capture file may contain multiple runs/data series. Display a list of runs and their start/end time and sample count, then select the newest one in `last_run` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = sqlite3.connect('data/waveform-raspi-ocxo-2day.sqlite3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run 000: 2020-04-01 14:00:25 - 2020-04-01 15:09:31 (  1:09:05.846,   4197664sp)\n",
      "Run 001: 2020-04-02 11:56:41 - 2020-04-02 11:57:59 (  0:01:18.544,     79552sp)\n",
      "Run 002: 2020-04-02 12:03:51 - 2020-04-03 14:12:18 (1 day, 2:08:27.618,  95262592sp)\n",
      "Run 003: 2020-04-03 14:12:48 - 2020-04-06 10:33:25 (2 days, 20:20:36.644, 249113600sp)\n"
     ]
    }
   ],
   "source": [
    "for run_id, start, end, count in db.execute('SELECT run_id, MIN(rx_ts), MAX(rx_ts), COUNT(*) FROM measurements GROUP BY run_id'):\n",
    "    foo = lambda x: datetime.datetime.fromtimestamp(x/1000)\n",
    "    start, end = foo(start), foo(end)\n",
    "    time_window = f'{start:%Y-%m-%d %H:%M:%S} - {end:%Y-%m-%d %H:%M:%S}'\n",
    "    print(f'Run {run_id:03d}: {time_window} ({str(end-start)[:-3]:>13}, {count*32:>9d}sp)')\n",
    "last_run, n_records = run_id, count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup analog parameters\n",
    "\n",
    "Setup parameters of analog capture hardware here. This is used to scale samples from ADC counts to analog voltages. Also setup sampling rate here. Nominal sampling rate is 1 ksps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_rate = 1000.0 * 48.6 / 48\n",
    "\n",
    "par = lambda *rs: 1/sum(1/r for r in rs) # resistor parallel calculation\n",
    "\n",
    "# Note: These are for the first prototype only!\n",
    "vmeas_source_impedance = 330e3\n",
    "vmeas_source_scale = 0.5\n",
    "\n",
    "vcc = 15.0\n",
    "vmeas_div_high = 27e3\n",
    "vmeas_div_low = par(4.7e3, 10e3)\n",
    "vmeas_div_voltage = vcc * vmeas_div_low / (vmeas_div_high + vmeas_div_low)\n",
    "vmeas_div_impedance = par(vmeas_div_high, vmeas_div_low)\n",
    "\n",
    "#vmeas_overall_factor = vmeas_div_impedance / (vmeas_source_impedance + vmeas_div_impedance)\n",
    "v0 = 1.5746\n",
    "v100 = 2.004\n",
    "vn100 = 1.1452\n",
    "\n",
    "adc_vcc = 3.3 # V\n",
    "adc_fullscale = 4095\n",
    "\n",
    "adc_val_to_voltage_factor = 1/adc_fullscale * adc_vcc\n",
    "\n",
    "adc_count_to_vmeas = lambda x: (x*adc_val_to_voltage_factor - v0) / (v100-v0) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load run data from sqlite3 capture file\n",
    "\n",
    "Load measurement data for the selected run and assemble a numpy array containing one continuous trace. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d5d47d8c9ff441febc6c31ce1b612abd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/7784800 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RMS voltage: 227.28577854695376\n"
     ]
    }
   ],
   "source": [
    "limit = n_records\n",
    "record_size = 32\n",
    "skip_dropped_sections = False\n",
    "\n",
    "data = np.zeros(limit*record_size)\n",
    "data[:] = np.nan\n",
    "\n",
    "last_seq = None\n",
    "write_index = 0\n",
    "for i, (seq, chunk) in tqdm(enumerate(db.execute(\n",
    "        'SELECT seq, data FROM measurements WHERE run_id = ? ORDER BY rx_ts LIMIT ? OFFSET ?',\n",
    "        (last_run, limit, n_records-limit))), total=n_records):\n",
    "    \n",
    "    if last_seq is None or seq == (last_seq + 1)%0x10000:\n",
    "        last_seq = seq\n",
    "        idx = write_index if skip_dropped_sections else i\n",
    "        data[idx*record_size:(idx+1)*record_size] = np.frombuffer(chunk, dtype='<H')\n",
    "        write_index += 1\n",
    "        \n",
    "    elif seq > last_seq:\n",
    "        last_seq = seq\n",
    "        # nans = np.empty((record_size,))\n",
    "        # nans[:] = np.nan\n",
    "        # data = np.append(data, nans) FIXME\n",
    "        \n",
    "data = (data * adc_val_to_voltage_factor - v0) / (v100-v0) * 100\n",
    "\n",
    "# https://stackoverflow.com/questions/6518811/interpolate-nan-values-in-a-numpy-array\n",
    "nan_helper = lambda y: (np.isnan(y), lambda z: z.nonzero()[0])\n",
    "\n",
    "# data rarely may contain NaNs where the capture script failed to read and acknowledge capture buffers from the sensor board fast enough.\n",
    "# For RMS calculation and overall FFT fill these NaNs with interpolated values from their neighbors.\n",
    "data_interp = np.copy(data)\n",
    "nans, x = nan_helper(data)\n",
    "data_interp[nans]= np.interp(x(nans), x(~nans), data[~nans])\n",
    "\n",
    "print('RMS voltage:', np.sqrt(np.mean(np.square(data_interp))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of skipped sample packets: 16\n",
      "Consecutive skipped packets: 1 pkt: 10 2 pkt: 3\n"
     ]
    }
   ],
   "source": [
    "import itertools\n",
    "skip_groups = [ len(list(group))//32 for val, group in itertools.groupby(nans) if val ]\n",
    "print('Number of skipped sample packets:', sum(skip_groups))\n",
    "print('Consecutive skipped packets:', ' '.join(f'{val} pkt: {len(list(group))}' for val, group in itertools.groupby(sorted(skip_groups))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show a preview of loaded data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Font family ['normal'] not found. Falling back to DejaVu Sans.\n",
      "findfont: Font family ['normal'] not found. Falling back to DejaVu Sans.\n",
      "findfont: Font family ['normal'] not found. Falling back to DejaVu Sans.\n"
     ]
    },
    {
     "data": {
      "application/pdf": "[encoded data, 114384 characters]\n",
      "image/png": "[encoded data, 220664 characters]\n",
      "text/plain": [
       "<Figure size 648x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, (top, bottom) = plt.subplots(2, figsize=(9,6))\n",
    "fig.tight_layout(pad=3, h_pad=1.8)\n",
    "\n",
    "range_start, range_len = -300, 60 # [s]\n",
    "\n",
    "data_slice = data[ int(range_start * sampling_rate) : int((range_start + range_len) * sampling_rate) ]\n",
    "\n",
    "top.grid()\n",
    "top.plot(np.linspace(0, range_len, int(range_len*sampling_rate)), data_slice, lw=1.0)\n",
    "top.set_xlim([range_len/2-0.25, range_len/2+0.25])\n",
    "mean = np.mean(data_interp)\n",
    "rms = np.sqrt(np.mean(np.square(data_interp - mean)))\n",
    "peak = np.max(np.abs(data_interp - mean))\n",
    "top.axhline(mean, color='red')\n",
    "bbox = {'facecolor': 'black', 'alpha': 0.8, 'pad': 2}\n",
    "top.text(0.02, 0.5, f'mean: {mean:.3f}', transform=top.transAxes, color='white', bbox=bbox, ha='left', va='center')\n",
    "top.text(0.98, 0.2, f'V_RMS: {rms:.3f}', transform=top.transAxes, color='white', bbox=bbox, ha='right')\n",
    "top.text(0.98, 0.1, f'V_Pk: {peak:.3f}', transform=top.transAxes, color='white', bbox=bbox, ha='right')\n",
    "top.text(0.5, 0.9, f'Run {run_id}', transform=top.transAxes, color='white', bbox=bbox, ha='center', fontweight='bold')\n",
    "\n",
    "bottom.grid()\n",
    "bottom.specgram(data_slice, Fs=sampling_rate)\n",
    "top.set_ylabel('U [V]')\n",
    "bottom.set_ylabel('F [Hz]')\n",
    "bottom.set_xlabel('t [s]')\n",
    "\n",
    "top.set_title('Voltage waveform')\n",
    "bottom.set_title('Voltage frequency spectrum')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate Short-Time Fourier Transform of capture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Window length: 256 sp, zero-padded to 256 sp\n",
      "STFT sampling rate: 7.91015625\n"
     ]
    }
   ],
   "source": [
    "def calc_stft(data, fs=sampling_rate, ff:'Hz nominal'=50.0):\n",
    "    analysis_periods = 10\n",
    "    window_len = 256 # fs * analysis_periods/ff\n",
    "    nfft_factor = 1\n",
    "    sigma = window_len/8 # samples\n",
    "\n",
    "    f, t, Zxx = signal.stft(data,\n",
    "                fs = fs,\n",
    "                window=('gaussian', sigma),\n",
    "                nperseg = window_len,\n",
    "                nfft = window_len * nfft_factor)\n",
    "    print(f'Window length: {window_len:.0f} sp, zero-padded to {window_len * nfft_factor:.0f} sp')\n",
    "    stft_output_sampling_rate = 1.0/(t[1] - t[0])\n",
    "    print('STFT sampling rate:', stft_output_sampling_rate)\n",
    "    return f, t, Zxx, stft_output_sampling_rate\n",
    "\n",
    "f, t, Zxx, stft_output_sampling_rate = calc_stft(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show a preview of STFT results\n",
    "\n",
    "Cut out our approximate frequency range of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/pdf": "[encoded data, 209752 characters]\n",
      "image/png": "[encoded data, 10468 characters]\n",
      "text/plain": [
       "<Figure size 648x216 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 3))\n",
    "fig.tight_layout(pad=2, h_pad=0.1)\n",
    "\n",
    "ax.pcolormesh(t[-200:-100], f[:250], np.abs(Zxx[:250,-200:-100]))\n",
    "ax.set_title(f\"Run {last_run}\", pad=-20, color='white')\n",
    "ax.grid()\n",
    "ax.set_ylabel('f [Hz]')\n",
    "ax.set_ylim([30, 75]) # Hz\n",
    "ax.set_xlabel('capture time t [s]')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Gasior and Gonzalez for precise frequency estimation\n",
    "\n",
    "Limit analysis to frequency range of interest. If automatic adaption to totally different frequency ranges\n",
    "(e.g. 400Hz) would be necessary, we could switch here based on configuration or a lookup of the STFT bin\n",
    "containing highest overall energy.\n",
    "\n",
    "As elaborated in the Gasior and Gonzalez Paper [1] the shape of the template function should match the expected peak shape.\n",
    "Peak shape is determined by the STFT window function. As Gasior and Gonzalez note, a gaussian is a very good fit for a steep gaussian window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54307c438a4f4db18224385b62c16abc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1946199 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def runner(args):\n",
    "    frame_f, frame_Z, le_t = args\n",
    "    # Template function. We use a gaussian here. This function needs to fit the window above.\n",
    "    def gauss(x, *p):\n",
    "        A, mu, sigma = p\n",
    "        return A*np.exp(-(x-mu)**2/(2.*sigma**2))\n",
    "\n",
    "    # Calculate initial values for curve fitting\n",
    "    f_start = frame_f[np.argmax(frame_Z)] # index of strongest bin index\n",
    "    A_start = np.max(frame_Z) # strongest bin value\n",
    "    p0 = [A_start, f_start, 1.]\n",
    "    try:\n",
    "        # Fit template to measurement data STFT ROI \n",
    "        coeff, var = optimize.curve_fit(gauss, frame_f, frame_Z, p0=p0)\n",
    "        _A, rv, _sigma, *_ = coeff # The measured frequency is the mean of the fitted gaussian\n",
    "        return rv\n",
    "\n",
    "    except Exception as e:\n",
    "        # Handle fit errors\n",
    "        return np.nan\n",
    "\n",
    "def gasior_gonzalez_fmeas(f, t, Zxx):\n",
    "    import multiprocessing\n",
    "\n",
    "    chunksize = 1000\n",
    "\n",
    "    n_f, n_t = Zxx.shape\n",
    "    # Frequency ROI\n",
    "    f_min, f_max = 30, 70 # Hz\n",
    "    # Indices of bins within ROI\n",
    "    bounds_f = slice(np.argmax(f > f_min), np.argmin(f < f_max))\n",
    "\n",
    "    # Initialize output array\n",
    "    f_mean = np.zeros(Zxx.shape[1])\n",
    "\n",
    "    jobs = {}\n",
    "    with multiprocessing.Pool(multiprocessing.cpu_count()) as pool, tqdm(total=Zxx.shape[1]-2) as tq:\n",
    "        # Iterate over STFT time slices\n",
    "        for le_t in range(0, Zxx.shape[1], chunksize):\n",
    "            # Cut out ROI and compute magnitude of complex fourier coefficients\n",
    "            jobs[le_t] = pool.map_async(runner, [\n",
    "                    (f[bounds_f], np.abs(Zxx[bounds_f, frame_t]), frame_t) for frame_t in range(le_t, min(Zxx.shape[1], le_t+chunksize))\n",
    "                ], callback=lambda _x: tq.update(chunksize))\n",
    "\n",
    "        pool.close()\n",
    "        for le_t, future in jobs.items():\n",
    "            f_mean[le_t:le_t+chunksize] = future.get()\n",
    "        pool.join()\n",
    "\n",
    "    # Cut off invalid values at fringes\n",
    "    return f_mean[1:-2], t[1:-2]\n",
    "\n",
    "f_mean, f_mean_t = gasior_gonzalez_fmeas(f, t, Zxx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Produce plots of measurement results\n",
    "\n",
    "### Plot results as time-series data\n",
    "Include measurements of mean, standard deviation and variance of measurement data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/pdf": "[encoded data, 53240 characters]\n",
      "image/png": "[encoded data, 54072 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "pdate = lambda s: dates.date2num(datetime.datetime.fromisoformat(s))\n",
    "td2num = lambda td: dates.date2num(start + datetime.timedelta(seconds=td))\n",
    "\n",
    "def plot_freq_trace(outfile, xlim=None, minor_locator=dates.HourLocator(interval=3), major_locator=dates.HourLocator(interval=12), smooth_sec=None):\n",
    "    fig, ax = plt.subplots(figsize=(6, 4), sharex=True)\n",
    "    fig.tight_layout(pad=3.5, h_pad=0, w_pad=1)\n",
    "\n",
    "    if smooth_sec is not None:\n",
    "        # smooth data by convolving with a blackman window\n",
    "        a = int(10 * smooth_sec)\n",
    "        a = a//2*2 + 1 # make a odd\n",
    "        w = np.blackman(a)\n",
    "        f_smooth = np.convolve(w/w.sum(), f_mean, mode='valid')\n",
    "        t_smooth = f_mean_t[a//2:-a//2+1]\n",
    "    else:\n",
    "        f_smooth = f_mean\n",
    "        t_smooth = f_mean_t\n",
    "    \n",
    "    ax.plot([ td2num(td) for td in t_smooth ], f_smooth, lw=1)\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "    \n",
    "    var = np.var(f_mean[~np.isnan(f_mean)][1:-1])\n",
    "    mean = np.mean(f_mean[~np.isnan(f_mean)][1:-1])\n",
    "    ax.text(0.5,  0.93, f'Run {run_id}, {time_window}', transform=ax.transAxes, ha='center', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    ax.text(0.05, 0.15, f'μ={mean:.3g} Hz', transform=ax.transAxes, ha='left', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    ax.text(0.05, 0.09, f'σ={np.sqrt(var) * 1e3:.3g} mHz', transform=ax.transAxes, ha='left', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    ax.text(0.05, 0.03, f'σ²={var * 1e3:.3g} mHz²', transform=ax.transAxes, ha='left', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "\n",
    "    # Indicate missing values\n",
    "    for i in np.where(np.isnan(f_mean))[0]:\n",
    "        ax.axvspan(td2num(t[i]), td2num(t[i+1]), color='lightblue')\n",
    "\n",
    "    ax.xaxis.set_major_locator(major_locator)\n",
    "    ax.xaxis.set_minor_locator(minor_locator)\n",
    "    formatter = dates.DateFormatter('%Y-%m-%d\\n%H:%M:%S')\n",
    "    ax.xaxis.set_major_formatter(formatter)\n",
    "    ax.set_xlabel('recording time t [hh:mm:ss]')\n",
    "    ax.grid(True, which='both')\n",
    "    if xlim is not None:\n",
    "        ax.set_xlim(xlim)\n",
    "    fig.savefig(f'fig_out/{outfile}.pdf')\n",
    "    None\n",
    "    \n",
    "plot_freq_trace('freq_meas_trace_24h',\n",
    "                xlim=[pdate('2020-04-03 15:00:00'), pdate('2020-04-05 21:00:00')],\n",
    "                smooth_sec=60*5,\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/pdf": "[encoded data, 47012 characters]\n",
      "image/png": "[encoded data, 36284 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_freq_trace('freq_meas_trace_2h_1',\n",
    "                xlim=[pdate('2020-04-03 23:00:00'), pdate('2020-04-04 01:00:00')],\n",
    "                smooth_sec=10,\n",
    "                minor_locator=dates.MinuteLocator(interval=10),\n",
    "                major_locator=dates.MinuteLocator(interval=30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/pdf": "[encoded data, 47028 characters]\n",
      "image/png": "[encoded data, 36960 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_freq_trace('freq_meas_trace_2h_2',\n",
    "                xlim=[pdate('2020-04-03 17:30'), pdate('2020-04-03 19:30:00')],\n",
    "                smooth_sec=10,\n",
    "                minor_locator=dates.MinuteLocator(interval=10),\n",
    "                major_locator=dates.MinuteLocator(interval=30))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot raw mains voltage spectrum\n",
    "\n",
    "First compute FFT of voltage, then smoothen and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_voltage_fft(data):\n",
    "    # Number of samplepoints\n",
    "    N = len(data)\n",
    "    # sample spacing\n",
    "    T = 1.0 / sampling_rate\n",
    "    x = np.linspace(0.0, N*T, N)\n",
    "    yf = np.absolute(scipy.fftpack.fft(data * sig.blackman(N)))**2\n",
    "    xf = np.linspace(0.0, 1.0/(2.0*T), N//2)\n",
    "\n",
    "    yf = 2.0/N * np.abs(yf[:N//2])\n",
    "\n",
    "    average_from = lambda val, start, average_width: np.hstack([val[:start], [ np.mean(val[i:i+average_width]) for i in range(start, len(val), average_width) ]])\n",
    "\n",
    "    average_width = 6\n",
    "    average_start = 20\n",
    "    yf = average_from(yf, average_start, average_width)\n",
    "    xf = average_from(xf, average_start, average_width)\n",
    "    yf = average_from(yf, 200, average_width)\n",
    "    xf = average_from(xf, 200, average_width)\n",
    "    return xf, yf\n",
    "#voltage_fft = compute_voltage_fft(data_interp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/pdf": "[encoded data, 105044 characters]\n",
      "image/png": "[encoded data, 21844 characters]\n",
      "text/plain": [
       "<Figure size 432x216 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def plot_voltage_spectrum(xf, yf):\n",
    "    fig, ax = plt.subplots(figsize=(6, 3))\n",
    "    fig.tight_layout(pad=2.0)\n",
    "\n",
    "    yf2 = np.copy(yf)\n",
    "    #chunk_size = 200\n",
    "    #for i in range(len(yf)-chunk_size, 1, -chunk_size):\n",
    "    #    yf2 = np.hstack([yf2[:i], sig.savgol_filter(yf2[i:], 23, 3)])\n",
    "\n",
    "    ax.loglog(xf, yf2, lw=1)\n",
    "    #ax.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _pos: f'{1/x:.1f}'))\n",
    "    ax.set_xlabel('f [Hz]')\n",
    "    ax.set_ylabel('Amplitude V [V]')\n",
    "    ax.grid()\n",
    "    ax.set_xlim([0.001, 500])\n",
    "    fig.subplots_adjust(bottom=0.2)\n",
    "\n",
    "    for le_f in (50, 150, 250, 350, 450):\n",
    "        ax.axvline(le_f, color=(1, 0.5, 0.5), zorder=-2)\n",
    "    ax.annotate('50 Hz', xy=(15, 0.1), xycoords='data', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    ax.text(0.5, 0.9, f'Run {run_id}, {time_window}', transform=ax.transAxes, ha='center', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    fig.savefig('fig_out/mains_voltage_spectrum.pdf')\n",
    "#plot_voltage_spectrum(*voltage_fft)\n",
    "plot_voltage_spectrum(*compute_voltage_fft(data_interp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot frequency spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Font family ['normal'] not found. Falling back to DejaVu Sans.\n"
     ]
    },
    {
     "data": {
      "application/pdf": "[encoded data, 85332 characters]\n",
      "image/png": "[encoded data, 29320 characters]\n",
      "text/plain": [
       "<Figure size 288x180 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def plot_fmeas_spectrum(data):\n",
    "    # Number of samplepoints\n",
    "    newcopy = np.copy(data)\n",
    "    nans, x = nan_helper(newcopy)\n",
    "    newcopy[nans]= np.interp(x(nans), x(~nans), newcopy[~nans])\n",
    "\n",
    "    N = len(newcopy)\n",
    "    # sample spacing\n",
    "    T = 1.0 / stft_output_sampling_rate\n",
    "    x = np.linspace(0.0, N*T, N)\n",
    "    yf = np.absolute(scipy.fftpack.fft(newcopy * sig.blackman(N)))**2\n",
    "    xf = np.linspace(0.0, stft_output_sampling_rate/2, N//2)\n",
    "\n",
    "    yf = 2.0/N * np.abs(yf[:N//2])\n",
    "\n",
    "    #chunk_size = 200\n",
    "    #for i in range(len(yf)-chunk_size, 1, -chunk_size):\n",
    "    #    yf = np.hstack([yf[:i], sig.savgol_filter(yf[i:], 23, 3)])\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(4,2.5))\n",
    "    ax.loglog(xf, yf, label='spectrum', lw=1, color='lightgray')\n",
    "    ax.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _pos: f'{1/x:.1f}'))\n",
    "    ax.set_xlabel('Period T [s]')\n",
    "    ax.set_ylabel('Power Spectral Density [Hz^2/Hz]')\n",
    "\n",
    "    intervals = [0.5, 1.0, 1.5, 2.0, 3.9, 6.3, 10, 60, 300, 450, 900, 1200, 1800, 2700, 3600]\n",
    "    for i, t in enumerate(intervals):\n",
    "        ax.axvline(1/t, color='red', alpha=0.3, zorder=-1, lw=1)\n",
    "        ax.annotate(f'{t} s', xy=(1/t, 1e-7), xytext=(-4, 5), xycoords='data', textcoords='offset pixels', rotation=90, fontsize=5)\n",
    "    \n",
    "    closest = [ np.abs(xf - 1/t).argmin() for t in intervals ]\n",
    "    w = lambda i: i//2+3\n",
    "    closest = np.array([ yf[max(0, i-w(i)//2):i+w(i)//2].argmax() + max(0, i-w(i)//2) for i in closest])\n",
    "    ax.scatter(1/np.array(intervals), yf[closest], zorder=100, marker='x', color='red', lw=1, s=4*4)\n",
    "    \n",
    "    #ax.text(1/60, 10,'60 s', ha='left')\n",
    "    ax.grid()\n",
    "    ax.set_xlim([1/(10*3600), 5])\n",
    "    ax.set_ylim([1e-7, 1e2])\n",
    "    ax.autoscale(False) # do not include noise illustration lines in autoscaling\n",
    "    ax.plot(xf[1:], 1e-5/xf[1:], label='$f^{-1}$ line', color='blue', lw=1)\n",
    "    #ax.plot(xf[1:], 1e-8/(xf[1:]**4), label='$f^{-3}$ line', color='orange', ls='--')\n",
    "    #ax.plot(xf[1:], np.tile(1e-6, len(xf)-1), label='noise floor', color='orange', ls='-.')\n",
    "    \n",
    "\n",
    "    #yf2 = np.copy(yf)\n",
    "    #chunk_size = 50\n",
    "    #for i in range(len(yf)-chunk_size, 1, -chunk_size):\n",
    "    #    yf2 = np.hstack([yf2[:i], sig.savgol_filter(yf2[i:], 23, 3)])\n",
    "    #chunk_size = 2000\n",
    "    #for i in range(len(yf2)-chunk_size, 1, -chunk_size):\n",
    "    #    yf2 = np.hstack([yf2[:i], sig.savgol_filter(yf2[i:], 511, 3)])\n",
    "    #ax.plot(xf, yf2)\n",
    "    spline_first = 4\n",
    "    foo = np.log(yf[spline_first:])\n",
    "    foo_w = np.tile(1, len(foo))\n",
    "    foo_w[np.isnan(foo)] = 0\n",
    "    foo[np.isnan(foo)] = 0\n",
    "    spl = scipy.interpolate.splrep(x=xf[spline_first:], y=foo, w=foo_w, t=np.logspace(-4.5, 0, 100) * (0.45 * stft_output_sampling_rate), task=-1)\n",
    "    ax.plot(xf[spline_first:], np.exp(scipy.interpolate.splev(xf[spline_first:], spl)), color='red', label='log-log spline\\ninterpolation', lw=1)\n",
    "    \n",
    "    ax.text(0.5, 0.9, f'Run {run_id}', transform=ax.transAxes, ha='center', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "\n",
    "    ax.legend(loc='upper right')\n",
    "    fig.tight_layout()\n",
    "    fig.savefig('fig_out/freq_meas_spectrum.pdf')\n",
    "    \n",
    "    return (xf[spline_first], xf[-1], len(xf[spline_first:])), spl\n",
    "psd_spl_x, psd_spl = plot_fmeas_spectrum(f_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export measurement data for modulation simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Invalid samples: 24 / 1946198 (12.3 ppm)\n"
     ]
    }
   ],
   "source": [
    "print(f'Invalid samples: {np.sum(np.isnan(f_mean))} / {len(f_mean)} ({np.sum(np.isnan(f_mean))/len(f_mean)*1e6:.1f} ppm)')\n",
    "\n",
    "with open(f'data/fmeas_export_ocxo_2day.bin', 'wb') as f:\n",
    "    for sample in f_mean:\n",
    "        if not np.isnan(sample):\n",
    "            f.write(struct.pack('<f', sample))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3.10/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n",
      "/usr/lib/python3.10/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n",
      "/usr/lib/python3.10/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n",
      "/usr/lib/python3.10/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n"
     ]
    },
    {
     "data": {
      "application/pdf": "[encoded data, 815704 characters]\n",
      "image/png": "[encoded data, 219388 characters]\n",
      "text/plain": [
       "<Figure size 864x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with open(f'grid_freq_psd_spl_{len(psd_spl[1])}pt.json', 'w') as f:\n",
    "    json.dump({'x_spec': psd_spl_x, 't': psd_spl[0].tolist(), 'c': psd_spl[1].tolist(), 'k': psd_spl[2] }, f)\n",
    "\n",
    "def generate_synthetic_noise(specfile='grid_freq_psd_spl_108pt.json'):\n",
    "    with open(specfile) as f:\n",
    "        d = json.load(f)\n",
    "        x = np.linspace(*d['x_spec'])\n",
    "        N = len(x)\n",
    "        psd_spl = d['t'], d['c'], d['k']\n",
    "        \n",
    "    noise = np.random.normal(size=N) * 2\n",
    "    spec = scipy.fftpack.fft(noise) **2\n",
    "\n",
    "    spec *= np.exp(scipy.interpolate.splev(x, psd_spl))\n",
    "\n",
    "    spec **= 1/2\n",
    "    \n",
    "    renoise = scipy.fftpack.ifft(spec)\n",
    "    return x, renoise\n",
    "\n",
    "def noise_sim(specfile='grid_freq_psd_spl_108pt.json'):\n",
    "    x, renoise = generate_synthetic_noise(specfile)\n",
    "    N = len(x)\n",
    "    respec = 2.0/N * np.absolute(scipy.fftpack.fft(renoise * np.blackman(N))) ** 2\n",
    "    \n",
    "    #xf = np.linspace(0, 10/2, N//2)\n",
    "    fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(5, figsize=(12, 15), gridspec_kw={'hspace': 0.5})\n",
    "    \n",
    "    ax1.loglog(x, respec, lw=1)\n",
    "    ax1.plot(x, np.exp(scipy.interpolate.splev(x, psd_spl)), color='red', label='log-log spline interpolation')\n",
    "    ax1.grid()\n",
    "    ax1.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _pos: f'{1/x:.1f}'))\n",
    "    ax1.set_xlabel('Period T [s]')\n",
    "    ax1.set_ylabel('Power Spectral Density\\n[$Hz^2/Hz$]')\n",
    "    ax1.set_title('Simulated Noise Spectrum')\n",
    "    def plot_sigs(ax, lims, tick_ivl, legend_loc=None, title=None):\n",
    "        ax.plot(renoise[slice(*lims)], label='Synthetic signal')\n",
    "        ax.plot(f_mean[slice(*lims)] - np.mean(f_mean[~np.isnan(f_mean)]), label='Reference signal')\n",
    "        ax.grid()\n",
    "        if legend_loc is not None:\n",
    "            ax.legend(loc=legend_loc)\n",
    "        ax.set_ylabel('Frequency deviation\\nΔf [Hz]')\n",
    "        ax.set_xlabel('Sampling time [s]')\n",
    "        ax.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _pos: f'{x/10:.0f}'))\n",
    "        ax.xaxis.set_major_locator(plt.MultipleLocator(tick_ivl * 10.0))\n",
    "        ax.set_xlim([0, lims[1]-lims[0]])\n",
    "        if title:\n",
    "            ax.set_title(title)\n",
    "    plot_sigs(ax2, [10000, 210000], 3600, title='Simulated Noise Waveform [5 hour]')\n",
    "    plot_sigs(ax3, [10000, 20000], 60, title='Simulated Noise Waveform [20 min]')\n",
    "    plot_sigs(ax4, [10000, 11000], 10, title='Simulated Noise Waveform [2 min]')\n",
    "    plot_sigs(ax5, [10000, 10100], 1, legend_loc='lower center', title='Simulated Noise Waveform [10 sec]')\n",
    "    \n",
    "    return fig\n",
    "    \n",
    "noise_sim().savefig('fig_out/simulated_noise_spectrum.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Window length: 256 sp, zero-padded to 256 sp\n",
      "STFT sampling rate: 7.91015625\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "455b4242e2eb4b4cb7f2a7fb06aa5985",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/28476 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.001270923391184572\n"
     ]
    },
    {
     "data": {
      "application/pdf": "[encoded data, 658244 characters]\n",
      "image/png": "[encoded data, 221892 characters]\n",
      "text/plain": [
       "<Figure size 864x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def do_artificial_noise_simulation(duration:'seconds' = 3600.0, sampling_rate=sampling_rate):\n",
    "    t_pad = 1000.0\n",
    "    offx = int(t_pad*stft_output_sampling_rate)\n",
    "    \n",
    "    _x, noise_freqs = generate_synthetic_noise()\n",
    "    noise_freqs = np.absolute(noise_freqs)[offx:][:int(duration * stft_output_sampling_rate) + 1]\n",
    "    x = np.linspace(0, duration, int(duration*sampling_rate))\n",
    "    noise_resampled = np.interp(x, np.linspace(0, len(noise_freqs)/stft_output_sampling_rate, len(noise_freqs)), noise_freqs)\n",
    "    \n",
    "    phase_acc = 0.0\n",
    "    out = np.zeros(len(noise_resampled))\n",
    "    for i, f in enumerate(noise_resampled):\n",
    "        phase_acc += 2*np.pi*(50.0 + f) / sampling_rate\n",
    "        out[i] = np.sin(phase_acc)\n",
    "        if phase_acc > 2*np.pi:\n",
    "            phase_acc -= 2*np.pi\n",
    "            \n",
    "    return out, noise_freqs\n",
    "\n",
    "def recalc_f(duration=3600.0):\n",
    "    new_data, orig_noise = do_artificial_noise_simulation(duration, sampling_rate)\n",
    "    f, t, Zxx, stft_output_sampling_rate = calc_stft(new_data, sampling_rate)\n",
    "    #plt.matshow(np.absolute(Zxx), aspect='auto')\n",
    "    f_mean, t = gasior_gonzalez_fmeas(f, t, Zxx)\n",
    "    return t, f_mean, orig_noise\n",
    "\n",
    "def feedback_plot(duration=3600):\n",
    "    fig, axs = plt.subplots(5, 1, figsize=(12, 15))\n",
    "    new_t, new_mean, orig_noise = recalc_f()\n",
    "    for ax, time_range in zip(axs.flatten(), (duration, 300, 30, 5)):\n",
    "        ax.plot(new_t, orig_noise[1:-1], label='original')\n",
    "        ax.plot(new_t, new_mean - 50, label='re-calculated')\n",
    "        ax.grid()\n",
    "        ax.set_xlim((duration/2-time_range/2, duration/2+time_range/2))\n",
    "        ax.set_ylabel('Frequency deviation\\nΔf [Hz]')\n",
    "    axs[-2].legend()\n",
    "    delta = orig_noise[1:-1] - (new_mean - 50)\n",
    "    print(np.sqrt(np.mean(np.square(delta))))\n",
    "    axs[-1].plot(new_t, delta)\n",
    "    axs[-1].set_xlabel('Sampling time [s]')\n",
    "    axs[-1].set_xlim((0, duration))\n",
    "    axs[-1].grid()\n",
    "    axs[-1].set_title('Difference')\n",
    "    axs[-1].set_ylabel('Frequency deviation\\nΔf [Hz]')\n",
    "    axs[0].set_title('Original and re-calculated signals')\n",
    "    fig.tight_layout()\n",
    "    return fig\n",
    "\n",
    "feedback_plot().savefig('fig_out/freq_meas_feedback.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "\n",
    "1. **Gasior, M. & Gonzalez, J.** Improving FFT frequency measurement resolution by parabolic and gaussian interpolation *CERN-AB-Note-2004-021, CERN-AB-Note-2004-021, 2004*"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jupytersucks-dsss",
   "language": "python",
   "name": "jupytersucks-dsss"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
