{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import struct\n",
    "\n",
    "import numpy as np\n",
    "from scipy import signal, optimize\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import rocof_test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "from IPython.display import set_matplotlib_formats\n",
    "%matplotlib widget\n",
    "#%matplotlib inline\n",
    "#set_matplotlib_formats('png', 'pdf')\n",
    "#font = {'family' : 'normal',\n",
    "#        'weight' : 'normal',\n",
    "#        'size'   : 10}\n",
    "#matplotlib.rc('font', **font)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 1000 # Hz\n",
    "ff = 50 # Hz\n",
    "duration = 60 # seconds\n",
    "# test_data = rocof_test_data.sample_waveform(rocof_test_data.test_close_interharmonics_and_flicker(),\n",
    "#                                             duration=20,\n",
    "#                                             sampling_rate=fs,\n",
    "#                                             frequency=ff)[0]\n",
    "# test_data = rocof_test_data.sample_waveform(rocof_test_data.gen_noise(fmin=10, amplitude=1),\n",
    "#                                             duration=20,\n",
    "#                                             sampling_rate=fs,\n",
    "#                                             frequency=ff)[0]\n",
    "\n",
    "test_data = []\n",
    "test_labels = [ fun.__name__.replace('test_', '') for fun in rocof_test_data.all_tests ] + ['tmp']\n",
    "for gen in rocof_test_data.all_tests:\n",
    "    test_data.append(rocof_test_data.sample_waveform(gen(),\n",
    "                                            duration=duration,\n",
    "                                            sampling_rate=fs,\n",
    "                                            frequency=ff)[0])\n",
    "test_data.append(rocof_test_data.sample_waveform(rocof_test_data.gen_chirp(49.8, 50.6, 0.40, dwell_time=0.60), duration=duration, sampling_rate=fs, frequency=ff)[0])\n",
    "# d = 10 # seconds\n",
    "# test_data = np.sin(2*np.pi * ff * np.linspace(0, d, int(d*fs)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "spr_fmt = f'{fs}Hz' if fs<1000 else f'{fs/1e3:f}'.rstrip('.0') + 'kHz'\n",
    "for label, data in zip(test_labels, test_data):\n",
    "    with open(f'rocof_test_data/rocof_test_{label}_{spr_fmt}.bin', 'wb') as f:\n",
    "        for sample in data:\n",
    "            f.write(struct.pack('<f', sample))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysis_periods = 10\n",
    "window_len = 256 # fs * analysis_periods/ff\n",
    "nfft_factor = 8\n",
    "sigma = window_len/8 # samples\n",
    "quantization_bits = 14\n",
    "\n",
    "ffts = []\n",
    "for item in test_data:\n",
    "    f, t, Zxx = signal.stft((item * (2**(quantization_bits-1) - 1)).round().astype(np.int16).astype(float),\n",
    "                fs = fs,\n",
    "                window=('gaussian', sigma),\n",
    "                nperseg = window_len,\n",
    "                nfft = window_len * nfft_factor)\n",
    "                #boundary = 'zeros')\n",
    "    ffts.append((f, t, Zxx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(129, 470)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Zxx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.90625"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1000/256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "429b30daffb04963ab07c34115ecb84b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(len(test_data), figsize=(8, 20), sharex=True)\n",
    "fig.tight_layout(pad=2, h_pad=0.1)\n",
    "\n",
    "for fft, ax, label in zip(test_data, ax.flatten(), test_labels):\n",
    "    ax.plot((item * (2**(quantization_bits-1) - 1)).round())\n",
    "    \n",
    "    ax.set_title(label, pad=-20, color='white', bbox=dict(boxstyle=\"square\", ec=(0,0,0,0), fc=(0,0,0,0.8)))\n",
    "    ax.grid()\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "ax.set_xlim([5000, 5200])\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c2ca60f3c6d489290a770d78195c4fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-27-31c82486a777>:6: MatplotlibDeprecationWarning: shading='flat' when X and Y have the same dimensions as C is deprecated since 3.3.  Either specify the corners of the quadrilaterals with X and Y, or pass shading='auto', 'nearest' or 'gouraud', or set rcParams['pcolor.shading'].  This will become an error two minor releases later.\n",
      "  ax.pcolormesh(t[1:], f[:250], np.abs(Zxx[:250,1:]))\n"
     ]
    }
   ],
   "source": [
    "fig, ax = plt.subplots(len(test_data), figsize=(8, 20), sharex=True)\n",
    "fig.tight_layout(pad=2, h_pad=0.1)\n",
    "\n",
    "for fft, ax, label in zip(ffts, ax.flatten(), test_labels):\n",
    "    f, t, Zxx = fft\n",
    "    ax.pcolormesh(t[1:], f[:250], np.abs(Zxx[:250,1:]))\n",
    "    ax.set_title(label, pad=-20, color='white')\n",
    "    ax.grid()\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "    ax.set_ylim([30, 75]) # Hz\n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  0.     ,   3.90625,   7.8125 ,  11.71875,  15.625  ,  19.53125,\n",
       "        23.4375 ,  27.34375,  31.25   ,  35.15625,  39.0625 ,  42.96875,\n",
       "        46.875  ,  50.78125,  54.6875 ,  58.59375,  62.5    ,  66.40625,\n",
       "        70.3125 ,  74.21875,  78.125  ,  82.03125,  85.9375 ,  89.84375,\n",
       "        93.75   ,  97.65625, 101.5625 , 105.46875, 109.375  , 113.28125,\n",
       "       117.1875 , 121.09375, 125.     , 128.90625, 132.8125 , 136.71875,\n",
       "       140.625  , 144.53125, 148.4375 , 152.34375, 156.25   , 160.15625,\n",
       "       164.0625 , 167.96875, 171.875  , 175.78125, 179.6875 , 183.59375,\n",
       "       187.5    , 191.40625, 195.3125 , 199.21875, 203.125  , 207.03125,\n",
       "       210.9375 , 214.84375, 218.75   , 222.65625, 226.5625 , 230.46875,\n",
       "       234.375  , 238.28125, 242.1875 , 246.09375, 250.     , 253.90625,\n",
       "       257.8125 , 261.71875, 265.625  , 269.53125, 273.4375 , 277.34375,\n",
       "       281.25   , 285.15625, 289.0625 , 292.96875, 296.875  , 300.78125,\n",
       "       304.6875 , 308.59375, 312.5    , 316.40625, 320.3125 , 324.21875,\n",
       "       328.125  , 332.03125, 335.9375 , 339.84375, 343.75   , 347.65625,\n",
       "       351.5625 , 355.46875, 359.375  , 363.28125, 367.1875 , 371.09375,\n",
       "       375.     , 378.90625, 382.8125 , 386.71875, 390.625  , 394.53125,\n",
       "       398.4375 , 402.34375, 406.25   , 410.15625, 414.0625 , 417.96875,\n",
       "       421.875  , 425.78125, 429.6875 , 433.59375, 437.5    , 441.40625,\n",
       "       445.3125 , 449.21875, 453.125  , 457.03125, 460.9375 , 464.84375,\n",
       "       468.75   , 472.65625, 476.5625 , 480.46875, 484.375  , 488.28125,\n",
       "       492.1875 , 496.09375, 500.     ])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b361be278748475bbe442f861a99418b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n",
      "/usr/lib/python3.9/site-packages/scipy/optimize/minpack.py:828: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n"
     ]
    }
   ],
   "source": [
    "fig, axs = plt.subplots(len(test_data)-1, figsize=(12, 15), sharex=True)\n",
    "axs = axs.flatten()\n",
    "\n",
    "for fft, label in zip(ffts, test_labels):\n",
    "    if label in ['noise_loud']: # custom test case, not part of upstream suite\n",
    "        continue\n",
    "    ax, *axs = axs\n",
    "        \n",
    "    f, f_t, Zxx = fft\n",
    "    \n",
    "    n_f, n_t = Zxx.shape\n",
    "    f_min, f_max = 30, 70 # Hz\n",
    "    bounds_f = slice(np.argmax(f > f_min), np.argmin(f < f_max))\n",
    "    \n",
    "    f_mean = np.zeros(Zxx.shape[1])\n",
    "    for t in range(1, Zxx.shape[1] - 1):\n",
    "        frame_f = f[bounds_f]\n",
    "        frame_step = frame_f[1] - frame_f[0]\n",
    "        time_step = f_t[1] - f_t[0]\n",
    "        frame_Z = np.abs(Zxx[bounds_f, t])\n",
    "        \n",
    "        def gauss(x, *p):\n",
    "            A, mu, sigma = p\n",
    "            return A*np.exp(-(x-mu)**2/(2.*sigma**2))\n",
    "        \n",
    "        f_start = frame_f[np.argmax(frame_Z)]\n",
    "        A_start = np.max(frame_Z)\n",
    "        p0 = [A_start, f_start, 1.]\n",
    "        try:\n",
    "            coeff, var = optimize.curve_fit(gauss, frame_f, frame_Z, p0=p0)\n",
    "            A, mu, sigma, *_ = coeff\n",
    "            f_mean[t] = mu\n",
    "        except RuntimeError:\n",
    "            f_mean[t] = np.nan\n",
    "    ax.plot(f_t[1:-1], f_mean[1:-1])\n",
    "    \n",
    "    ax.set_title(label, pad=-20, bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    ax.set_ylabel('f [Hz]')\n",
    "    ax.grid()\n",
    "    if not label in ['off_frequency', 'sweep_phase_steps']:\n",
    "        #ax.set_ylim([49.90, 50.10])\n",
    "        var = np.var(f_mean[1:-1])\n",
    "        ax.text(0.5, 0.1, f'σ²={var * 1e3:.3g} mHz²', transform=ax.transAxes, ha='center', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "        ax.text(0.5, 0.25, f'σ={np.sqrt(var) * 1e3:.3g} mHz', transform=ax.transAxes, ha='center', bbox=dict(fc='white', alpha=0.8, ec='none'))\n",
    "    else:\n",
    "        f_min, f_max = min(f_mean[1:-1]), max(f_mean[1:-1])\n",
    "        delta = f_max - f_min\n",
    "        ax.set_ylim(f_min - delta * 0.1, f_max + delta * 0.3)\n",
    "        \n",
    "ax.set_xlabel('simulation time t [s]')\n",
    "fig.tight_layout(pad=2.2, h_pad=0, w_pad=1)\n",
    "fig.savefig('fig_out/freq_meas_rocof_reference.pdf')\n",
    "None"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ma-thesis-env",
   "language": "python",
   "name": "ma-thesis-env"
  },
  "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
