{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Golem Workshop Kiten 2018 - Data Processing\n",
    "\n",
    "This is a notebook for data processing of the Golem operation domain scan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# At first, load some basic libraries.\n",
    "\n",
    "import numpy as np\n",
    "from urllib.request import urlopen\n",
    "import shutil\n",
    "import os \n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib import cm\n",
    "\n",
    "# Plot settings, can be changed to globally influence the appearance\n",
    "\n",
    "plt.rc('text', usetex=True)\n",
    "plt.rc('font', family='sans-serif')\n",
    "plt.rc('font', size=22)          # controls default text sizes\n",
    "plt.rc('axes', titlesize=22)     # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=20)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=20)    # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=20)    # fontsize of the tick labels\n",
    "plt.rc('legend', fontsize=22)    # legend fontsize\n",
    "plt.rc('figure', titlesize=22)  # fontsize of the figure title"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Golem Database Access\n",
    "\n",
    "Thanks to Golem staff for the golem_data function.\n",
    "\n",
    "This should not be changed :)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def golem_data(shot, diagn):\n",
    "    \"\"\"\n",
    "     Simple interface for GOLEM database\n",
    "      Use:\n",
    "\tobj = golem_data(10011, 'loop_voltage')\n",
    "\tplot(obj.tvec, obj.data)\n",
    "\n",
    "\td - object containing all availible informations\n",
    "\td.tvec  - time vector Nx1\n",
    "\td.data - data matrix  NxM\n",
    "\n",
    "      Example:\n",
    "      from golem_data import golem_data\n",
    "      from matplotlib.pyplot import *\n",
    "      obj1 = golem_data(10689 , 'electron_temperature')\n",
    "      obj2 = golem_data(10689 , 'spectrometr:temperature')\n",
    "      plot(obj1.tvec,obj1.data,  label=obj1.labels)  %\n",
    "      errorbar(obj2.tvec, obj2.data, xerr=obj2.tvec_err, yerr=[obj2.data_err[:,0], obj2.data_err[:,1]],  label=obj2.labels )\n",
    "      xlabel(obj2.ax_labels[0])\n",
    "      ylabel(obj2.ax_labels[1])\n",
    "      legend([obj1.name, obj2.name])\n",
    "      axis([obj1.plasma_start, obj1.plasma_end, 0, None])\n",
    "      title(obj1.info)\n",
    "      show()\n",
    "    \"\"\"\n",
    "      \n",
    "    remote_file = \"http://golem.fjfi.cvut.cz/utils/data/\"+str(shot)+\"/\"+diagn+'.npz'\n",
    "    \n",
    "    gfile = np.DataSource().open(remote_file, mode='rb')  # remote file\n",
    "    try:\n",
    "        d =  dict(np.load(gfile, encoding='bytes', fix_imports=True))\n",
    "    except IOError:\n",
    "        raise IOError('Missing diagnostic ' + str(diagn) + ' in shot ' + str(shot))\n",
    "    \n",
    "    if not 'tvec' in d and 't_start' in d:\n",
    "        d['tvec'] = np.linspace( d.pop('t_start'), d.pop('t_end'), len(d['data']) )\n",
    "    \n",
    "    try:\n",
    "        if 'scale' in d:\n",
    "            d['data'] = np.double(d['data']) * d.pop('scale')  # rescale data \n",
    "    except:\n",
    "        pass\n",
    "    return Data( **d )\n",
    "\n",
    "\n",
    "class Data(object):\n",
    "    \"\"\" Simple data handling object\"\"\"\n",
    "    def __init__(self, **kwargs):\n",
    "        self.__dict__.update(kwargs)\n",
    "    def __repr__(self):\n",
    "        return \"Data object, keys:\\n\" + str(self.__dict__.keys())\n",
    "    def __getitem__(self, key):\n",
    "        return getattr(self, key)\n",
    "\t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some Useful Scripts\n",
    "\n",
    "### Validate\n",
    "Excludes bad shots from a shot list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Validate(shot_list, bad_shots):\n",
    "    valid = [];\n",
    "    for shot_no in shot_list:\n",
    "        try:\n",
    "            if not(shot_no in bad_shots): \n",
    "                b_v = golem_data(shot_no, 'breakdown_voltage').data\n",
    "                valid.append(shot_no)\n",
    "        except Exception:\n",
    "            pass\n",
    "    return valid;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GetValAtTime\n",
    "\n",
    "Takes a diagnostic obtained by `golem_data` and interpolates its value for a specific time in the shot.\n",
    "\n",
    "Calling with `plot_progress = True` will show the interpolation on the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def GetValAtTime(diag, time, plot_progress = False):\n",
    "    f = interp1d(diag.tvec, diag.data)\n",
    "    # change to true if you want to see the plots\n",
    "    if plot_progress:\n",
    "        fig = plt.figure()\n",
    "        plt.plot(diag.tvec, diag.data, linewidth=2)\n",
    "        t = np.linspace(min(diag.tvec), max(diag.tvec), 50)\n",
    "        plt.plot(t, f(t), linewidth=1)\n",
    "        plt.plot([time, time], [min(diag.data), max(diag.data)], '--')\n",
    "        plt.plot(time, f(time), 'o')\n",
    "        plt.xlabel('$t$ [s]')\n",
    "        plt.ylabel('$B_\\mathrm{t}$ [T]')\n",
    "        plt.grid(True, axis='both')\n",
    "\n",
    "        plt.show()\n",
    "        plt.close()\n",
    "    return f(time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load\n",
    "\n",
    "Loads all specified diagnostics for a specified shot number.\n",
    "\n",
    "Any newly requested diagnostics should be added here.\n",
    "\n",
    "**To use the diagnostic in the scan, one point should be selected (e.g. by interpolation for a specific time).** Toroidal field is now treated like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Diagnostics should be added here (according to the sample fro)\n",
    "def Load(shot_no):\n",
    "    data_dict = {\n",
    "        'shot_no': shot_no,\n",
    "        'pressure_request': golem_data(shot_no, 'pressure_request').data,\n",
    "        'pressure': golem_data(shot_no, 'pressure').data,\n",
    "        'ucd': golem_data(shot_no, 'ucd').data,\n",
    "        'electron_temperature': golem_data(shot_no, 'electron_temperature_max').data,\n",
    "        'breakdown_voltage': golem_data(shot_no, 'breakdown_voltage').data,\n",
    "        'toroidal_field': GetValAtTime(golem_data(shot_no, 'toroidal_field'), \n",
    "                                       golem_data(shot_no, 'plasma_start').data)\n",
    "        }\n",
    "    return data_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GetScan1, GetScan2\n",
    "\n",
    "Sorts the data for further processing. Poorly written, should be somewhat more refined (to allow for arbitrary N-dimensional scans)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def GetScan1(data_list, scan_var, diag_var):\n",
    "    d_l = sorted(data_list, key=lambda d: d[scan_var])\n",
    "    return {\n",
    "            'x': [d[scan_var] for d in d_l],\n",
    "            'val': [d[diag_var] for d in d_l],\n",
    "            }\n",
    "\n",
    "def GetScan2(data_list, scan_var, diag_var):\n",
    "    d_l = sorted(data_list, key=lambda d: (d[scan_var[0]], d[scan_var[1]]))\n",
    "    return {\n",
    "            'x': [d[scan_var[0]] for d in d_l],\n",
    "            'y': [d[scan_var[1]] for d in d_l],\n",
    "            'val': [d[diag_var] for d in d_l],\n",
    "            }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ScanPlot2D\n",
    "\n",
    "Plots a 2D scan (scaled colormap)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ScanPlot2D(golem_scan, scan_id, scan_description):\n",
    "    fig = plt.figure(figsize=(8, 6), dpi=80)\n",
    "\n",
    "    plt.tricontourf(golem_scan[scan_id]['x'], golem_scan[scan_id]['y'], \n",
    "                    golem_scan[scan_id]['val'])\n",
    "    plt.plot(golem_scan[scan_id]['x'], golem_scan[scan_id]['y'], 'ko', ms=3)\n",
    "    plt.xlabel('$p$ [mPa]')\n",
    "    plt.ylabel('$U_\\mathrm{CD}$ [V]')\n",
    "    c = plt.colorbar()\n",
    "    c.set_label(scan_description)\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ScanPlotLines\n",
    "\n",
    "Plots a 2D scan as a series of lines grouped by one of scan variables (can be selected by setting `group_by` to either `x` or `y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ScanPlotLines(golem_scan, group_by, scan_id, scan_description):\n",
    "    fig = plt.figure(figsize=(8, 6), dpi=80)\n",
    "\n",
    "    x_axis = 'x' if group_by == 'y' else 'y';\n",
    "    \n",
    "    x_label = '$p$ [mPa]' if x_axis == 'x' else '$U_\\mathrm{CD}$ [V]'\n",
    "    lines_label = 'mPa' if x_axis == 'y' else 'V'\n",
    "    \n",
    "    unique = set(golem_scan[scan_id][group_by])\n",
    "    unique = sorted(list(unique))\n",
    "    for u_v in unique:\n",
    "        data_x = []\n",
    "        data_y = []\n",
    "        for i in range(len(golem_scan[scan_id][x_axis])):\n",
    "            if golem_scan[scan_id][group_by][i] == u_v:\n",
    "                data_x.append(golem_scan[scan_id][x_axis][i])\n",
    "                data_y.append(golem_scan[scan_id]['val'][i])\n",
    "        \n",
    "        plt.plot(data_x, data_y, 'o--', label = '{0} {1}'.format(u_v, lines_label))\n",
    "    \n",
    "    plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)\n",
    "    plt.xlabel(x_label)\n",
    "    plt.ylabel(scan_description)\n",
    "    plt.grid()\n",
    "    \n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Load\n",
    "\n",
    "Let's load the data from selected shots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Kiten scan\n",
    "shot_list = range(27454, 27491, 1)\n",
    "bad_shots = [27481, 27482, 27483]\n",
    "\n",
    "# Second scan (Prague 28/06)\n",
    "shot_list = [27617, 27618, 27619, 27620, 27621, 27622, 27623, 27624, 27625, 27626, 27627, 27628, 27629, 27630, \n",
    "             27633, 27634, 27635, 27636, 27637, 27638, 27639, 27640, 27641, 27642, 27643, 27644, 27645, 27646, \n",
    "             27647, 27648, 27649, 27650, 27651, 27652, 27653, 27654, 27655, 27656, 27657, 27658, 27659, 27660, \n",
    "             27661, 27662, 27663, 27664, 27665, 27666, 27667, 27668, 27669, 27670, 27671, 27672, 27673, 27674]\n",
    "bad_shots = []\n",
    "\n",
    "# Do the loading\n",
    "valid_shots = Validate(shot_list, bad_shots)\n",
    "data_list = [Load(s_no) for s_no in valid_shots]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Processing\n",
    "\n",
    "Three scans variables (`['breakdown_voltage', 'electron_temperature', 'toroidal_field']`) are selected and a scans in `pressure` and `ucd` (current drive voltage) are generated.\n",
    "\n",
    "A second scan set is generated for requested pressure instead of real pressure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "scan_vars = ['breakdown_voltage', 'electron_temperature', 'toroidal_field']\n",
    "\n",
    "# the main scan\n",
    "golem_scan = {};\n",
    "for s_id in scan_vars:\n",
    "    golem_scan[s_id] = GetScan2(data_list, ['pressure', 'ucd'], s_id)\n",
    "    \n",
    "# the auxiliary scan\n",
    "golem_scan_request = {};\n",
    "for s_id in scan_vars:\n",
    "    golem_scan_request[s_id] = GetScan2(data_list, ['pressure_request', 'ucd'], s_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Breakdown voltage\n",
    "We should see some minima on the Paschen curves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ScanPlotLines(golem_scan, 'y', 'breakdown_voltage', '$U_\\mathrm{b}$ [V]');\n",
    "fig.savefig(\"grouped-p-U_bd.pdf\", bbox_inches=\"tight\");\n",
    "fig.savefig(\"grouped-p-U_bd.png\", bbox_inches=\"tight\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2D plots\n",
    "\n",
    "* Breakdown voltage\n",
    "* Electron temperature\n",
    "* Toroidal field at the beginning of the discharge (plasma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ScanPlot2D(golem_scan, 'breakdown_voltage', '$U_\\mathrm{b}$ [V]');\n",
    "fig.savefig(\"scan2d-p-U_cd-U_bd.pdf\", bbox_inches=\"tight\");\n",
    "fig.savefig(\"scan2d-p-U_cd-U_bd.png\", bbox_inches=\"tight\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig = ScanPlot2D(golem_scan, 'electron_temperature', '$T_\\mathrm{e}$ [eV]');\n",
    "fig.savefig(\"scan2d-p-U_cd-T_e.pdf\", bbox_inches=\"tight\");\n",
    "fig.savefig(\"scan2d-p-U_cd-T_e.png\", bbox_inches=\"tight\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ScanPlot2D(golem_scan, 'toroidal_field', '$B_\\mathrm{t}$ [T]');\n",
    "fig.savefig(\"scan2d-p-U_cd-B_t.pdf\", bbox_inches=\"tight\");\n",
    "fig.savefig(\"scan2d-p-U_cd-B_t.png\", bbox_inches=\"tight\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Toroidal field \n",
    "\n",
    "Grouped by the requested pressure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ScanPlotLines(golem_scan_request, 'x', 'toroidal_field', '$B_\\mathrm{t}$ [T]');\n",
    "fig.savefig(\"grouped-p-B_t.pdf\", bbox_inches=\"tight\");\n",
    "fig.savefig(\"grouped-p-B_t.png\", bbox_inches=\"tight\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}