{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac7ff5f0",
   "metadata": {},
   "source": [
    "# Analysis of DRP signals, analysis of particle transport\n",
    "\n",
    "The following script shows the basic data aquision from **D**ouble **R**ake **P**robe\n",
    "\n",
    "To get more detailed information about **D**ouble **R**ake **P**robe visit its [wikipage](http://golem.fjfi.cvut.cz/wiki/Diagnostics/ParticleFlux/DoubleRakeProbe/index).\n",
    "\n",
    "DRP consists of two rakes of Langmuir probes that can be operated in two basic modes.  \n",
    "The first one is floating potential mode ($V_\\mathrm{float}$), where each pin is insulated  \n",
    "and we can measure an approximation of plasma potential.  \n",
    "\n",
    "The other mode is ion saturation current ($I_\\mathrm{sat}$), where -100 V voltage is  \n",
    "applied to each pin. This voltage repels electrons so one can measure the plasma density.\n",
    "\n",
    "To measure transport using DRP one has to measure 4 floating potential ($V_\\mathrm{fl}$)  \n",
    "signals and 2 ion saturation current ($I_\\mathrm{sat}$) signals, however, quality $I_\\mathrm{sat}$  \n",
    "measurements may turn out to be challenging to perform, so in the following  \n",
    "script, only $V_\\mathrm{float}$ measurements are considered."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1806b5b",
   "metadata": {},
   "source": [
    "### Needed Python modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf34bef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "try:\n",
    "    import xrscipy.signal as dsp\n",
    "except:\n",
    "    print('xrscipy.signal module not found')\n",
    "# import xarray_dsp.wavelets as wave\n",
    "\n",
    "# for better readibility of spectrograms\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "\n",
    "# to keep the attributes\n",
    "xr.set_options(keep_attrs=True)\n",
    "\n",
    "# the data is downloaded in the directory/folder\n",
    "# from where this script was started \n",
    "ds = np.DataSource('/tmp')\n",
    "\n",
    "#inline plots\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f681704",
   "metadata": {},
   "source": [
    "## Download (and save the data)\n",
    "One might need to modify `ulr_base` variable in `download_data` function and the channel names in `channel_list`.\n",
    "\n",
    "\n",
    "The correct `url_base` depends on the DAS used and the `channel_list`   \n",
    "represents the names of particular signals.  \n",
    "It is recommended to save the data to your computer if runnig the scipts locally,  \n",
    "as it saves time and makes your analysis independent on internet connection.\n",
    "\n",
    "Eventually the data is converted into xarray.DataArray structure for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca57b87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test shot_no, the followin configuration should give the right data\n",
    "# shot_no = 36443\n",
    "# test channel_list\n",
    "# channel_list = ['DRP-R1', 'DRP-R2','DRP-R3','DRP-R4','DRP-R5','DRP-R6',\n",
    "#               'DRP-L1', 'DRP-L2' ,'DRP-L3' ,'DRP-L4' ,'DRP-L5']\n",
    "# test url_base\n",
    "# \"http://golem.fjfi.cvut.cz/shots/{}/Diagnostics/DoubleRakeProbe/\".format(shot_no)\n",
    "\n",
    "# a little helper function\n",
    "def strip_off_data(base_url: str,data_name: str):\n",
    "    \"\"\" This function is used to get data like parameters\n",
    "    or comments from golem database\"\"\"\n",
    "    url = base_url + data_name\n",
    "    file = ds.open(url)\n",
    "    parameters = np.loadtxt(file)\n",
    "    return parameters\n",
    "\n",
    "def check_plasma_existence(shot_no: int, no_plasma_override: bool=False):\n",
    "    \"\"\"Checks, whether there was plasma, if there was not plasma\n",
    "    raises RuntimeError. It is possible to override no_plasma and\n",
    "    proceed in the code execution.\"\"\"\n",
    "#     base_url = \"http://golem.fjfi.cvut.cz/shots/{}/Diagnostics/BasicDiagnostics/Results/\"\n",
    "    base_url = f\"http://golem.fjfi.cvut.cz/shots/{shot_no}/Diagnostics/PlasmaDetection/Results/\"    \n",
    "    # check whether there was plasma\n",
    "    \n",
    "    b_plasma = strip_off_data(base_url,\"b_plasma\")\n",
    "    if not b_plasma == 1.0:\n",
    "        if not no_plasma_override:\n",
    "            raise RuntimeError(\"There was no plasma.\")\n",
    "    \n",
    "            \n",
    "\n",
    "\n",
    "def download_data(shot_no: int, channel_list: list, no_plasma_override: bool=False):\n",
    "    \"\"\"Downloads the data and returns a tuple\n",
    "    of time numpy.ndarray and data numpy.ndarray\"\"\"\n",
    "    url_base = f\"http://golem.fjfi.cvut.cz/shots/{shot_no}/Diagnostics/DoubleRakeProbe/\"\n",
    "    \n",
    "    check_plasma_existence(shot_no, no_plasma_override)\n",
    "    \n",
    "    data_all = np.ones(shape=(40960,len(channel_list)))\n",
    "    i = 0\n",
    "    \n",
    "    \n",
    "\n",
    "    for channel in channel_list:\n",
    "        # open the file\n",
    "        file = ds.open(url_base+\"{}.csv\".format(channel))\n",
    "        # transform to numpy format\n",
    "        np_file = np.loadtxt(file, delimiter=',')\n",
    "        # add to the complete array\n",
    "        data_all[:,i] = np_file[:,1]\n",
    "        # saves the time from the diag. output\n",
    "        if i == 0:\n",
    "            time = np_file[:,0] * 10e2 # to have time in ms\n",
    "        i += 1\n",
    "        \n",
    "        # the 10e2 coefficient is given by a voltage divider\n",
    "    return time, data_all * 10e2\n",
    "\n",
    "# list of channel names in the database\n",
    "channel_list = ['DRP-R1', 'DRP-R2','DRP-R3','DRP-R4','DRP-R5','DRP-R6',\n",
    "               'DRP-L1', 'DRP-L2' ,'DRP-L3' ,'DRP-L4' ,'DRP-L5','DRP-L6']\n",
    "# shot is set to 0\n",
    "# i.e. the last shot\n",
    "# attention! when running the script locally\n",
    "# leaving the 0 will not update the downloaded data to the current shot\n",
    "# because it first looks wheather the data has already been downloaded\n",
    "\n",
    "shot_no = 44201\n",
    "\n",
    "\n",
    "time, data = download_data(shot_no, channel_list ,False)\n",
    "\n",
    "# creation of xarray.DataArray\n",
    "# it is with offset\n",
    "data_all_off = xr.DataArray(data,\n",
    "                            coords={'time':time,\n",
    "                                    'signal':channel_list},\n",
    "                            dims=['time','signal']\n",
    "#                             ,attrs = {'shot_no'}\n",
    "                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe47eed",
   "metadata": {},
   "source": [
    "### Download the time of the plasma start and the end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dca31d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_plasma_start_end(shot_no: int, no_plasma_override: bool=False):\n",
    "    \"\"\"Downloads the plasma start and end if there was plasma.\"\"\"\n",
    "    # the following should not change in principle, but\n",
    "    # if your are having problems with downloading the data,\n",
    "    # check the url as adviced above\n",
    "    base_url = f\"http://golem.fjfi.cvut.cz/shots/{shot_no}/Diagnostics/PlasmaDetection/Results/\"\n",
    "    \n",
    "    # checks whether there was plasma\n",
    "    # and if yes, get the plasma start and end\n",
    "    # this check is done here again because it is\n",
    "    # meaningless to try to download plasma_start and plasma_end\n",
    "    # if there was no plasma,\n",
    "    check_plasma_existence(shot_no, no_plasma_override)\n",
    "    \n",
    "    t_start = strip_off_data(base_url,\"t_plasma_start\")\n",
    "    t_end = strip_off_data(base_url,\"t_plasma_end\")\n",
    "    return t_start, t_end\n",
    "\n",
    "# add the plasma start and end as a attribute of \n",
    "# all_data xarray.DataArray\n",
    "t_start, t_end = get_plasma_start_end(shot_no,False)\n",
    "\n",
    "\n",
    "data_all_off.attrs = {\"plasma_start\":t_start, \"plasma_end\":t_end}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbcfc3e5",
   "metadata": {},
   "source": [
    "### Subtraction of offset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d274c1a9",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# time window of period before the plasma start\n",
    "# where we calculate the offset we subtract\n",
    "time_off = slice(data_all_off.attrs['plasma_start']-2.5,data_all_off.attrs['plasma_start']-1)\n",
    "\n",
    "# all the downloaded data without offset\n",
    "data_all = data_all_off - data_all_off.sel(time=time_off).mean(dim='time')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34e66573",
   "metadata": {},
   "source": [
    "### Visual check of the signals "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e451e447",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_window = slice(data_all.attrs['plasma_start']-0.5,data_all.attrs['plasma_end']+1)\n",
    "# the following list just to sort the pin coordinates such that\n",
    "# it is shown as it is physically in the vessel\n",
    "signal_order = ['DRP-R1', 'DRP-L1','DRP-R2','DRP-L2','DRP-R3','DRP-L3',\n",
    "               'DRP-R4', 'DRP-L4' ,'DRP-R5' ,'DRP-L5' ,'DRP-R6', 'DRP-L6']\n",
    "\n",
    "# here is the part, where the plasma is is selected\n",
    "sel_plasma = data_all.sel(time=time_window,signal=signal_order)\n",
    "\n",
    "# plot and plot setup\n",
    "all_plot = sel_plasma.plot.line(x=\"time\", col='signal', col_wrap=2,\n",
    "                                figsize=(12,12),\n",
    "#                                 subplot_kws={'gridspec_kw:{'wspace':0.1}}\n",
    "                               )\n",
    "\n",
    "fontsize = 13\n",
    "# axes labels\n",
    "all_plot.set_xlabels('time [ms]',fontsize=fontsize)\n",
    "all_plot.set_ylabels('$V_\\mathrm{float}$ [V]',fontsize=fontsize)\n",
    "\n",
    "\n",
    "# make the figure nicer\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8615188e",
   "metadata": {},
   "source": [
    "### Spectrograms of the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8e07e51",
   "metadata": {},
   "source": [
    "we keep using the signal order and time window\n",
    "from the cell above → so we continue to work\n",
    "with the `sel_plasma` variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27213caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrogram = dsp.spectrogram(sel_plasma, nperseg=512,dim='time')\n",
    "\n",
    "spct_plot = spectrogram.plot(x='time',col='signal',col_wrap=2,\n",
    "                             norm=LogNorm(),\n",
    "                             figsize=(12,12),\n",
    "                             cbar_kwargs={'label':r'_[-]'}\n",
    "                            )\n",
    "\n",
    "fontsize = 13\n",
    "# axes labels\n",
    "spct_plot.set_xlabels('time [ms]',fontsize=fontsize)\n",
    "spct_plot.set_ylabels('$f$ [kHz]',fontsize=fontsize)\n",
    "\n",
    "plt.ylim(0,75)\n",
    "# save the figure with the name expected by DoubleRakeProbe.sh\n",
    "plt.savefig('icon-fig.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd79a03",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}