{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to xemc3\n",
    "\n",
    "xemc3 is a library for reading the output fron EMC3 simulations into the xarray format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import xemc3\n",
    "\n",
    "# Matplotlib setup\n",
    "import setup_plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quick introduction to xarray\n",
    "\n",
    "This is just to get a quick introduction of the structure of the xarray data type.\n",
    "\n",
    "In the cell below we generate a xarray with dimensions $(3,3)$ for variable $x$ with coordinates $(10,20)$ and $y$ with coordinates $(1,2,3)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = xr.DataArray(\n",
    "    np.random.rand(2, 3), dims=(\"x\", \"y\"), coords={\"x\": [10, 20], \"y\": [1, 2, 3]}\n",
    ")\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data.values\n",
    "\n",
    "Returns the wrapped numpy array, in this case the return value of `np.random.rand(2, 3)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data.dims\n",
    "\n",
    "Returns the name of the dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.dims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data.coords\n",
    "Returns the coordinates for all axis directions with coordinate names and datatype of the coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.coords"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data.attrs\n",
    "\n",
    "Returns other attributes in form of a dictionary with you can easily add by generating a new value associated with a new key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.attrs[\"units\"] = \"m\"\n",
    "data.attrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More on xarray:\n",
    "* [Documentation](https://xarray.pydata.org/en/stable/index.html)\n",
    "* [Tutorials](https://xarray-contrib.github.io/xarray-tutorial/index.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EMC3 data\n",
    "\n",
    "The prerequisite for this example to work is to have downloaded the file emc3_example.nc and have the libraries specified in this script installed in your enviroment. We recommend using netCDF4 for opening .nc files. The emc3_example.nc can be found and downloaded here: https://gitlab.mpcdf.mpg.de/dave/xemc3-data given that you have acces.\n",
    "\n",
    "The path specified in the string in the cell below is where you have stored the emc3_example.nc locally on your computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use local helper function to get some data\n",
    "from get_data import load_example_data\n",
    "\n",
    "ds = load_example_data()\n",
    "# If you want to use your own data use something like\n",
    "# ds = xemc3.load.all(\"path/to/mydata/\")\n",
    "# or if you have converted it already to a netcdf file\n",
    "# ds = xr.open_dataset(\"path/to/mydata.nc\")\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### dataset (ds) explanation\n",
    "\n",
    "When running the codeline ds on the last line of a cell you get an overview of what the xarray object consist of.\n",
    "\n",
    "#### ds.coords['R_bounds']\n",
    "\n",
    "`R_bounds` represents the coordinates of the vertices at the gridcells in the radial direction in the $xy$-plane. with $R = \\sqrt{x^2 + y^2}$.\n",
    "\n",
    "#### ds.coords['z_bounds']\n",
    "`z_bounds` represents the coordinates of the vertices of the gridcells in the $z$-direction.\n",
    "\n",
    "#### ds.coords['phi_bounds']\n",
    "`phi_bounds` represents the coordinates of the vertices of the gridcells in the $\\phi$-direction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# The shape is 6 dimensional, as we include 3 dimensions for the vertices\n",
    "ds.coords[\"R_bounds\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that units are m - xemc3 prefers SI, with only eV as exception\n",
    "ds.coords[\"z_bounds\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Toroidal slice\n",
    "A toroidal slice is defined as the grid of $(R,z)$-values at a fixed angle $\\phi$. The values of the $\\phi$-angles used in the W7X grid can be found in the paragraph below and demonstrated in the next cell.\n",
    "\n",
    "### ds.coords['phi_bounds']\n",
    "Running the cell below gives you an array of the $\\phi$ angles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.coords[\"phi_bounds\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ds.emc3.plot_rz(key, phi = $\\phi$)\n",
    "\n",
    "The key is given as a string, `None` can be passed as a key if you want to plot the mesh. An example is the angle `phi` $= \\phi$ which is the angle given in radians as float. \n",
    "\n",
    "For this particular example(.nc file) the floats of the angle $\\phi$ can be found in the dictionary defined by ds.coords['phi_bounds'] which has 2 dimensions; one for either side of the cell for a given angle $\\phi$. There are 37 different values for $\\phi$ since W7-X has a five-fold symmetry which is divided in two up-down symmetric parts and we use a resolution $1^{\\circ}$.\n",
    "\n",
    "In the cells below are some examples of the parameter electron temperature $T_e$ plotted in toroidal slices for phi index $n_{\\phi} = [0,18,35]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "# the parameter can be plotted by a one-liner\n",
    "plt.figure()\n",
    "ds.emc3.plot_rz(\"Te\", phi=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for several angles and control\n",
    "import ipywidgets as widgets\n",
    "\n",
    "fig = plt.figure()\n",
    "\n",
    "\n",
    "def plot_Te(ip):\n",
    "    fig.clear()\n",
    "    ds.emc3.plot_rz(\"Te\", phi=ip)\n",
    "\n",
    "\n",
    "ip = widgets.FloatSlider(min=0, max=np.pi / 5, value=0, step=0.01)\n",
    "widgets.interact(plot_Te, ip=ip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simplifying of grid structure\n",
    "\n",
    "The parameter values are defined in each grid cell, but the center or the mean of the vertices of the gridcell: $\\mathbf{r}_{param} = \\langle \\mathbf{r}_{vertex} \\rangle$. A simplified analogy is the centerpoint of a 3D cube.\n",
    "\n",
    "You can get the center coordinates `ds.<direction>_bounds.mean(dim=(\"delta_r\", \"delta_theta\", \"delta_phi\"), ignore_missing=True)` by averaging over the `delta_*` dimensions of the variable, as that contains the cell extend in each direction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "R = ds.R_bounds.mean(dim=(\"delta_r\", \"delta_theta\", \"delta_phi\"))\n",
    "z = ds.z_bounds.mean(dim=(\"delta_r\", \"delta_theta\", \"delta_phi\"))\n",
    "phi = ds.phi_bounds.mean(dim=\"delta_phi\")\n",
    "x = R * np.cos(phi)\n",
    "y = R * np.sin(phi)\n",
    "x, ds.Te"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use of NaN values in xEMC3\n",
    "\n",
    "Not all gridcells have a defined parameter value attached to it. This is mostly the outer and inner region of the grid where the values of many parameter has been ignored because in the specific regions the quantity is not evolved.\n",
    "This is illustrated in the above plot example of the electron temperature $T_e$. In the cell below you can see how large a fraction of the total number of gridpoints the mesh for the electron temperature that has NaN as a value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_nans_Te = np.sum(np.isnan(ds.Te)).data\n",
    "print(\"How many nans in Te field?\", n_nans_Te)\n",
    "print(\n",
    "    f\"Fraction of nans with respect to gridcells: {n_nans_Te / ds.Te.size * 100:.2f} %\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grid\n",
    "\n",
    "In the cell below there is an interactive plot of the grid. You can use the slider to iterate through all toroidal slices(all $\\phi$ angles)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "\n",
    "fig = plt.figure()\n",
    "\n",
    "\n",
    "def plot_emc3grid(ip):\n",
    "    fig.clear()\n",
    "    ds.emc3.plot_rz(None, phi=ip)\n",
    "\n",
    "\n",
    "ip = ipywidgets.FloatSlider(min=0, max=np.pi / 5, value=0, step=0.01)\n",
    "ipywidgets.interact(plot_emc3grid, ip=ip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grid with boundaries\n",
    "\n",
    "Interactive plot of the grid, here you can use the ipywidget slider to iterate through all toroidal slices,\n",
    "the rmin and rmax to set the boundaries in r direction, and the zmin and zmax to set the boundaries in z direction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "ds.emc3.plot_rz(\"Te\", phi=np.pi / 5, Rmin=5.8, Rmax=6.1, zmin=0, zmax=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Periodic boundary conditions for plotting\n",
    "\n",
    "Naturally the data does not have periodic boundary conditions, which means that the last dataframe would be equal to the first. In the case of emc3 data the periodicity is in the theta direction. For plotting the dimension of the theta grid is increased by one and set to the first values in the theta direction. This is for tha plot to \"complete the orbit\" in the theta direction for it to be closed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "\n",
    "\n",
    "def plot_Te_zoomed(ip):\n",
    "    fig.clear()\n",
    "    c = plt.pcolormesh(\n",
    "        ds.emc3[\"R_corners\"][:, :, ip],\n",
    "        ds.emc3[\"z_corners\"][:, :, ip],\n",
    "        ds.Te[:, :, ip],\n",
    "        cmap=plt.cm.jet,\n",
    "        shading=\"auto\",\n",
    "    )\n",
    "    plt.colorbar(c)\n",
    "\n",
    "\n",
    "phislider = ipywidgets.IntSlider(min=0, max=35)\n",
    "ipywidgets.interact(plot_Te_zoomed, ip=phislider)"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}