{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1><center>IAGA Summer School 2019</center></h1>\n",
    "\n",
    "<h1><center>K Index Calculation</center></h1>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import notebook dependencies\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "sys.path.append('..')\n",
    "from src import mag_lib as mag\n",
    "\n",
    "# Set data paths\n",
    "ESK_2003_PATH = os.path.abspath('../data/external/ESK_2003/')\n",
    "K_INDS_PATH = os.path.abspath('../data/external/k_inds/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Calculating K-indices for a single observatory\n",
    "\n",
    "The K-index is a global geomagnetic activity index devised by Julius Bartels in 1938 to give a simple measure of the degree of geomagnetic disturbance during each 3-hour (UT) interval. It uses data from ground-based magnetometers to assign a number in the range 0-9 to interval, with K=0 indicating very little geomagnetic activity and K=9 representing an extreme geomagnetic storm. The K-index was introduced at the time of photographic recording, when magnetograms recorded variations in the horizontal geomagnetic field elements declination (D) and horizontal intensity (H), and in the vertical intensity (Z). To derive a K-index an observer would fit, by eye, a quiet-time curve to the records of D and H and measure the range (maximum-minimum) of the deviation of the recording from the curve. The K-index was then assigned according to a conversion table, with the greatest range in D and H 'winning'. The north component (X) may be used instead of H, and the east component (Y) instead of D (X and Y will be used in the examples below). The vertical component Z is not used because it is liable to contamination by local induced currents.\n",
    "\n",
    "The conversion from range in nanoteslas to index is quasi-logarithmic. The conversion table varies with latitude in an attempt to normalise the K-index distribution for observatories at different latitudes. The table for Eskdalemuir is shown below.\n",
    "\n",
    "| K | 0 | 1 |2|3|4|5|6|7|8|9|\n",
    "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n",
    "| Lower bound (nT) | 0 | 8 |15|30|60|105|180|300|500|750|\n",
    "\n",
    "This means that, for instance, K=2 if the measured range is in the interval \\[15, 30\\] nT.\n",
    "\n",
    "There was a VERY long debate in IAGA Division V about algorithms that could adequately reproduce the K-indices that an experienced observer would assign. The algorithms and code approved by IAGA are available at the International Service for Geomagnetic Indices: http://isgi.unistra.fr/softwares.php. (Also http://isgi.unistra.fr/what_are_kindices.php)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example\n",
    "\n",
    "In the following cells, we **illustrate** the process. A simple approach is to assume the so-called regular daily variation $S_R$ is made up of 24h, 12h and 8h signals (and, possibly, higher harmonics). A Fourier analysis can be used to investigate this. The functions in the cell below calculate Fourier coefficients from a data sample, and then synthesise a Fourier series using the coefficients.\n",
    "\n",
    "Some days this simple approach seems to work well, on others it's obviously wrong. Think about another approach you might take.\n",
    "\n",
    "We will then attempt to calculate K-indices for the day selected by computing the Fourier series up the the number of harmonics selected by subtracting the synthetic harmonic signal from the data, then calculating 3-hr ranges and converting these into the corresponding K-index. The functions to do are also in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fourier(v, nhar):\n",
    "    npts    = len(v)\n",
    "    f       = 2.0/npts\n",
    "    t       = np.linspace(0, npts, npts, endpoint=False)*2*np.pi/npts\n",
    "    vmn     = np.mean(v)\n",
    "    v       = v - vmn\n",
    "    cofs    = [0]*(nhar+1)\n",
    "    cofs[0] = (vmn,0)\n",
    "    for i in range(1,nhar+1):\n",
    "        c, s    = np.cos(i*t), np.sin(i*t)\n",
    "        cofs[i] = (np.dot(v,c)*f, np.dot(v,s)*f)\n",
    "    return (cofs)\n",
    "\n",
    "def fourier_synth(cofs, npts):\n",
    "    nt  = len(cofs)\n",
    "    syn = np.zeros(npts)\n",
    "    t  = np.linspace(0, npts, npts, endpoint=False)*2*np.pi/npts\n",
    "    for n in range(1, nt):\n",
    "        for j in range(npts):\n",
    "            syn[j] += cofs[n][0]*np.cos(n*t[j]) + cofs[n][1]*np.sin(n*t[j])\n",
    "    return (syn)\n",
    "\n",
    "def K_plot(d, synd, day, E):\n",
    "    fig, ax = plt.subplots(2, 1, figsize=(15, 10))\n",
    "    t = np.linspace(0, 1440, 1440, endpoint= False)/60\n",
    "    ax[0].plot(t, d)\n",
    "    ax[0].plot(t, synd, color='red')\n",
    "    ax[0].set_xlabel('UT (h)', fontsize=14)\n",
    "    ax[0].set_ylabel(f'{E} (nT)', fontsize=14)\n",
    "    ax[0].xaxis.set_ticks(np.arange(0, 27, 3))\n",
    "    ax[0].set_title(f'K-index estimation: {OBSY} {day}, ({E}-range: {np.ptp(d):.1f}nT)', fontsize=20)\n",
    "    ax[1].plot(t, d-synd)\n",
    "    ax[1].set_xlabel('UT (h)', fontsize=14)\n",
    "    ax[1].set_ylabel(f'{E} residuals (nT)', fontsize=14)\n",
    "    ax[1].xaxis.set_ticks(np.arange(0, 27, 3))\n",
    "    fig.tight_layout()\n",
    "\n",
    "def K_calc(d, synd, Kb):\n",
    "    tmp = np.ptp((d-synd).reshape(8,180), axis=1)\n",
    "    return(list(np.digitize(tmp, bins=list(Kb.values()), right=False)-1))\n",
    "\n",
    "def K_load(obscode, year, k_inds_path=K_INDS_PATH):\n",
    "    kfile  = year + '.' + obscode\n",
    "    filepath = os.path.join(k_inds_path, obscode, kfile)\n",
    "    df = pd.read_csv(filepath, skiprows=0, header=None, delim_whitespace=True,  \n",
    "                     parse_dates=[[2,1,0]], index_col=0)\n",
    "    df = df.drop(3, axis=1)\n",
    "    df.index.name='Date'\n",
    "    return(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, load in (X, Y, Z) one-minute data from Eskdalemuir for 2003 into a pandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OBSY = 'ESK'\n",
    "year ='2003'\n",
    "npts = 1440 # Number of minutes in a day\n",
    "df = mag.load_year('esk', 2003, ESK_2003_PATH)\n",
    "df.columns = [col.strip(OBSY) for col in df.columns]\n",
    "df = df.drop(['DOY', 'F'], axis=1)\n",
    "ann_means = df[[f'{i}' for i in 'XYZ']].resample('1Y').mean()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### >> USER INPUT HERE: Choose a day (2003-mm-dd) and the number of harmonics to calculate\n",
    "#### Return here to investigate different days (note that only the file for 2003 is included in the data directory for this example).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "day  = '2003-01-01'\n",
    "nhar = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First convert the D variations to units of nT using the annual mean value of H. Then calculate the daily means of D and H."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xmn  = df[day].X.mean()\n",
    "ymn  = df[day].Y.mean()\n",
    "xmn, ymn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now analyse the X data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x     = df[day].X.values - xmn\n",
    "xcofs = fourier(x, nhar)\n",
    "synx  = fourier_synth(xcofs, npts)\n",
    "K_plot(x, synx, day, 'X')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Repeat for Y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y     = df[day].Y.values - ymn\n",
    "ycofs = fourier(y, nhar)\n",
    "syny  = fourier_synth(ycofs, npts)\n",
    "K_plot(y, syny, day, 'Y')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the K-indices individually for the X and Y data, and also the final K index (the largest of the X and Y indices)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_lb   = [0, 8, 15, 30, 60, 105, 180, 300, 500, 750]\n",
    "K_dict = dict(zip([f'{OBSY}_K{i}' for i in range (10)], K_lb))\n",
    "Kx = K_calc(x, synx, K_dict)\n",
    "Ky = K_calc(y, syny, K_dict)\n",
    "Kf =[max(Kx[i], Ky[i]) for i in range(8)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the official 2003 observatory values for comparison and display the first few days."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfk   = K_load(OBSY.lower(), year)\n",
    "dfk.columns = ['00','03','06','09','12','15','18','21']\n",
    "K_obs = list(dfk.loc[day])\n",
    "dfk.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a data frame containing all K values and print them out to compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = pd.DataFrame({'K_x': Kx, 'K_y': Ky, 'K_final': Kf, 'K_obs': K_obs}).transpose()\n",
    "K.columns = ['0','3','6','9','12','15','18','21']\n",
    "\n",
    "print(\"Comparison of our K-indices with the official values\")\n",
    "K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot histograms of the observed K-indices for each 3-hour period during 2003 in separate subplots\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfk.hist(figsize=(12,12), bins=[i for i in range(11)], sharey=True, align='left', rwidth=0.8)\n",
    "plt.suptitle('ESK 2003: Distribution of K-indices for each 3-hour interval')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plotted side by side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7,7))\n",
    "plt.hist(dfk.values, bins=[i for i in range(11)], align='left')\n",
    "plt.legend(dfk.columns)\n",
    "plt.ylabel('Number of 3-hour intervals')\n",
    "plt.xlabel('K')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and stacked to give the distribution for the whole year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7,7))\n",
    "plt.hist(dfk.values, bins=[i for i in range(11)], stacked=True, align='left', rwidth=0.8)\n",
    "plt.legend(dfk.columns)\n",
    "plt.ylabel('Number of 3-hour intervals')\n",
    "plt.xlabel('K')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also compute a daily sum of the K-indices for the 2003 file, and list days with high and low summed values. Note that this summation is not really appropriate because the K-index is quasi-logarithmic, however, this is a common simple measure of quiet and disturbed days. (These might be interesting days for you to look at.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfk['Ksum'] = dfk.sum(axis=1)\n",
    "Ksort = dfk.sort_values('Ksum')\n",
    "print('Quiet days: \\n\\n', Ksort.head(10), '\\n\\n')\n",
    "print('Disturbed days: \\n\\n', Ksort.tail(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Note on the Fast Fourier Transform\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the examples above we computed Fourier coefficients in the 'traditional' way, so that if $F(t)$ is a Fourier series representation of $f(t)$, then,\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "F(t) &= A_o+\\sum_{n=1}^N A_n \\cos\\left(\\frac{2\\pi nt}{T}\\right)+B_n \\sin\\left(\\frac{2\\pi nt}{T}\\right)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where $T$ is the fundamental period of $F(t)$. The $A_n$ and $B_n$ are estimated by\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "A_o&=\\frac{1}{T}\\int_0^T f(t) dt\\\\\n",
    "A_n&=\\frac{2}{T}\\int_0^T f(t)\\cos\\left(\\frac{2\\pi nt}{T}\\right) dt\\\\\n",
    "B_n&=\\frac{2}{T}\\int_0^T f(t)\\sin\\left(\\frac{2\\pi nt}{T}\\right) dt\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "With $N$ samples of digital data, the integral for $A_n$ may be replaced by the summation\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "A_n&=\\frac{2}{T}\\sum_{j=0}^{N-1} f_j\\cos\\left(\\frac{2\\pi nj\\Delta t}{T}\\right) \\Delta t\\\\\n",
    "&=\\frac{2}{N}\\sum_{j=0}^{N-1} f_j\\cos\\left(\\frac{2\\pi nj}{N}\\right)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where the sampling interval $\\Delta t$ is given by $T = N \\Delta t$ and $f_j = f(j \\Delta t)$. A similar expression applies for the $B_n$, and these are the coefficients returned by the function _fourier_ above.\n",
    "\n",
    "The fast Fourier transform (FFT) offers a computationally efficient means of finding the Fourier coefficients. The conventions for the FFT and its inverse (IFFT) vary from package to package. In the _scipy.fftpack_ package, the FFT of a sequence $x_n$ of length $N$ is defined as\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "y_k&=\\sum_{n=0}^{N-1} x_n\\exp\\left(-\\frac{2\\pi i\\thinspace kn}{N}\\right)\\\\\n",
    "&=\\sum_{n=0}^{N-1} x_n\\left(\\cos\\left(\\frac{2\\pi \\thinspace kn}{N}\\right)-i\\sin\\left(\\frac{2\\pi \\thinspace kn}{N}\\right)\\right)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "with the inverse defined as,\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "x_n&=\\frac{1}{N}\\sum_{k=0}^{N-1} y_k\\exp\\left(\\frac{2\\pi i\\thinspace kn}{N}\\right)\\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The _scipy_ documentation is a little inconsistent here because it explains the order of the $y_n$ as being $y_1,y_2, \\dots y_{N/2-1}$ as corresponding to increasing positive frequency and $y_{N/2}, y_{N/2+1}, \\dots y_{N-1}$ as ordered by decreasing negative frequency, for $N$ even. (See: https://docs.scipy.org/doc/scipy/reference/tutorial/fftpack.html.)\n",
    "\n",
    "The interpretation is that if $y_k=a_k+ib_k$ then will have (for $N$ even), $y_{N-k} = a_k-ib_k$ and so\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "a_k&=\\frac{1}{2}\\text{Re}\\left(y_k+y_{N-k}\\right)\\\\\n",
    "b_k&=\\frac{1}{2}\\text{Im}\\left(y_k-y_{N-k}\\right)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "\n",
    "and so we expect the relationship to the digitised Fourier series coefficients returned by the function _fourier_ defined above to be,\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "A_k&=\\phantom{-}\\frac{1}{N}\\text{Re}\\left(a_k+a_{N-k}\\right)\\\\\n",
    "B_k&=-\\frac{1}{N}\\text{Im}\\left(b_k-b_{N-k}\\right)\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.fftpack import fft\n",
    "\n",
    "# Compute the fft of the x data\n",
    "xfft = fft(x)\n",
    "\n",
    "# Compare results for the 24-hour component\n",
    "k    = 1\n",
    "print('Fourier coefficients: \\n', f'A1 = {xcofs[1][0]} \\n', f'B1 = {xcofs[1][1]} \\n')\n",
    "print('scipy FFT outputs: \\n', f'a1 = {np.real(xfft[k]+xfft[npts-k])/npts} \\n', \\\n",
    "      f'b1 = {-np.imag(xfft[k]-xfft[npts-k])/npts} \\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "Menvielle, M. et al. (1995) ‘Computer production of K indices: review and comparison of methods’, Geophysical Journal International. Oxford University Press, 123(3), pp. 866–886. doi: 10.1111/j.1365-246X.1995.tb06895.x."
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}