{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cloudy-puzzle",
   "metadata": {},
   "source": [
    "# Geophysical Indices\n",
    "Geophysical indices are standard parameters that charaterize the global geomagnetic conditions and activity at a particular time.  These have standard definitions and are available from a variety of sources.  We will use the files hosted by [GFZ Potsdam](https://www.gfz-potsdam.de/en/kp-index/). The FTP site hosts yearly ASCII files that include the Kp, ap, and Ap indices, as well as the International Sunspot Number and the F10.7 Solar Radio Flux."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "rocky-evaluation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from ftplib import FTP\n",
    "import datetime as dt\n",
    "import numpy as np\n",
    "import h5py\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cosmetic-adult",
   "metadata": {},
   "source": [
    "## Download a file\n",
    "Yearly files are avialable from the [Potsdam FTP site](ftp://ftp.gfz-potsdam.de/pub/home/obs/Kp_ap_Ap_SN_F107).  These can be downloaded with the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "loving-antibody",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'221-You have transferred 61564 bytes in 1 files.\\n221-Total traffic for this session was 63431 bytes in 1 transfers.\\n221-Thank you for using the FTP service on ftp.\\n221 Goodbye.'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "year = 2015\n",
    "\n",
    "source_site = 'ftp.gfz-potsdam.de'\n",
    "ftp = FTP(source_site)\n",
    "ftp.login()\n",
    "filename = 'Kp_ap_Ap_SN_F107_{:4}.txt'.format(year)\n",
    "\n",
    "with open(os.path.join(filename), 'wb') as lf:\n",
    "    ftp.retrbinary('RETR /pub/home/obs/Kp_ap_Ap_SN_F107/{}'.format(filename), lf.write)\n",
    "ftp.quit()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comfortable-singapore",
   "metadata": {},
   "source": [
    "## Read in a file\n",
    "These ASCII files are fairly simple and consist of whitespace delimited values that can be read in with the numpy loadtxt function.  From the file header:"
   ]
  },
  {
   "cell_type": "raw",
   "id": "above-external",
   "metadata": {},
   "source": [
    "# Short file description (for a detailed file description, see Kp_ap_Ap_SN_F107_format.txt):\n",
    "# 40 header lines, all starting with #\n",
    "# ASCII, blank separated and fixed length, missing data indicated by -1.000 for Kp, -1 for ap and SN, -1.0 for F10.7\n",
    "# YYYY MM DD is date of UT day, days is days since 1932-01-01 00:00 UT to start of UT day, days_m is days since 1932-01-01 00:00 UT to midday of UT day\n",
    "# BSR is Bartels solar rotation number, dB is day within BSR\n",
    "# Kp1 to Kp8 (Kp for the eight eighth of the UT day), ap1 to ap8 (ap for the eight eighth of the UT day), Ap, SN, F10.7obs, F10.7adj\n",
    "# D indicates if the Kp and SN values are definitive or preliminary. D=0: Kp and SN preliminary; D=1: Kp definitive, SN preliminary; D=2 Kp and SN definitive\n",
    "#\n",
    "#\n",
    "# The format for each line is (i stands for integer, f for float):\n",
    "#iii ii ii iiiii fffff.f iiii ii ff.fff ff.fff ff.fff ff.fff ff.fff ff.fff ff.fff ff.fff iiii iiii iiii iiii iiii iiii iiii iiii  iiii iii ffffff.f ffffff.f i\n",
    "#\n",
    "# The parameters in each line are:\n",
    "#YYY MM DD  days  days_m  Bsr dB    Kp1    Kp2    Kp3    Kp4    Kp5    Kp6    Kp7    Kp8  ap1  ap2  ap3  ap4  ap5  ap6  ap7  ap8    Ap  SN F10.7obs F10.7adj D\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "hungarian-census",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.loadtxt(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "alien-taste",
   "metadata": {},
   "source": [
    "### Retrive indices for a particular time\n",
    "To retrieve indices for a particular date/time, simpily select the line from the data array that corresponds to the day in question.  For the parameters that depend on time of day (Kp and ap), you can select the correct value by dermining which 3-hour block the time lies within."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "polished-raise",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bsr 2476.0\n",
      "dB 15.0\n",
      "Kp 3.333\n",
      "ap 18.0\n",
      "Ap 9.0\n",
      "SN 76.0\n",
      "F10.7obs 153.0\n",
      "F10.7adj 148.9\n"
     ]
    }
   ],
   "source": [
    "targtime = dt.datetime(2015,2,7,11,4,0)\n",
    "\n",
    "# Find which line corresponds to \n",
    "lidx = np.argwhere((data[:,0]==targtime.year) & (data[:,1]==targtime.month) & (data[:,2]==targtime.day)).flatten()[0]\n",
    "# Select that line from the data array\n",
    "gpi_values = data[lidx]\n",
    "\n",
    "# Find which 3-hour indices should be used\n",
    "hidx = targtime.hour // 3\n",
    "\n",
    "# For convenience, save in a dictionary\n",
    "gpi_idx = {'Bsr':5, 'dB':6, 'Kp':7+hidx, 'ap':15+hidx, 'Ap':23, 'SN':24, 'F10.7obs':25, 'F10.7adj':26}\n",
    "gpi_out = {k:gpi_values[i] for k,i in gpi_idx.items()}\n",
    "\n",
    "for k, v in gpi_out.items():\n",
    "    print(k, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pending-imaging",
   "metadata": {},
   "source": [
    "### Create arrays of indices for a range of times\n",
    "Similarly, the indices can be extracted over a range of times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "critical-demand",
   "metadata": {},
   "outputs": [],
   "source": [
    "starttime = dt.datetime(2015,2,1)\n",
    "endtime = dt.datetime(2015,3,1)\n",
    "\n",
    "time_array = np.array([starttime+dt.timedelta(hours=h) for h in np.arange((endtime-starttime).total_seconds()/3600)])\n",
    "\n",
    "gpi_dates = np.array([dt.date(int(d[0]),int(d[1]),int(d[2])) for d in data])\n",
    "gpi_idx = {'Bsr':5, 'dB':6, 'Kp':7, 'ap':15, 'Ap':23, 'SN':24, 'F10.7obs':25, 'F10.7adj':26}\n",
    "gpi_out = {k:list() for k in gpi_idx.keys()}\n",
    "\n",
    "for t in time_array:\n",
    "    lidx = np.argwhere(t.date()==gpi_dates).flatten()[0]\n",
    "    \n",
    "    # Find which 3-hour indices should be used\n",
    "    hidx = t.hour // 3\n",
    "\n",
    "    for k, i in gpi_idx.items():\n",
    "        if k=='Kp' or k=='ap':\n",
    "            gpi_out[k].append(data[lidx,i+hidx])\n",
    "        else:\n",
    "            gpi_out[k].append(data[lidx,i])\n",
    "\n",
    "gpi_out = {k:np.array(v) for k,v in gpi_out.items()}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "floral-genesis",
   "metadata": {},
   "source": [
    "## Plot Geophsical Indices\n",
    "Create a simple plot of the Kp, ap, Ap, sunspot Number (SN), and adjusted F10.7 Solar Radio Flux (F10.7adj)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ranking-tiger",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x110a856d0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52452 characters]\n",
      "text/plain": [
       "<Figure size 1080x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(15,5))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "for k in ['Kp', 'ap', 'Ap', 'SN', 'F10.7adj']:\n",
    "    ax.plot(time_array, gpi_out[k], label=k)\n",
    "    \n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pursuant-charm",
   "metadata": {},
   "source": [
    "## Save Output File\n",
    "Save indices to an output drivers hdf5 file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "married-lighter",
   "metadata": {},
   "outputs": [],
   "source": [
    "utime = np.array([(t-dt.datetime.utcfromtimestamp(0)).total_seconds() for t in time_array])\n",
    "\n",
    "with h5py.File('sample_drivers.h5', 'w') as h5:\n",
    "    h5.create_dataset('UnixTime', data=utime)\n",
    "    for k, v in gpi_out.items():\n",
    "        h5.create_dataset(k, data=v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "internal-buddy",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
