{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate at\n",
    "\n",
    "Some examples of the evaluate_at functionality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import xemc3\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Matplotlib setup\n",
    "import setup_plt"
   ]
  },
  {
   "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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate along a line of sight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nx = 1000\n",
    "mapped = ds.emc3.evaluate_at_xyz(\n",
    "    np.linspace(4.8, 6.0, nx),\n",
    "    np.linspace(-0.1, 0.1, nx),\n",
    "    np.linspace(-0.5, 0.5, nx),\n",
    "    \"ne\",\n",
    "    periodicity=5,\n",
    "    updownsym=True,\n",
    "    delta_phi=np.pi / 1800,\n",
    ")\n",
    "# Add a coordinate\n",
    "mapped.coords[\"dim_0\"] = np.linspace(4.8, 6.0, nx)\n",
    "mapped.dim_0.attrs = dict(units=\"m\", long_name=\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "mapped[\"ne\"].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate along a 2D array\n",
    "\n",
    "Plotting an $R\\times z$ plane can be done with `ds.emc3.plot_rz(key, phi)` this allows to plot abitrary cuts of the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate an abitrary 2D slice\n",
    "x = xr.DataArray(np.linspace(0, 7, 100), name=\"x\", dims=\"x\", attrs=dict(units=\"m\"))\n",
    "y = xr.DataArray(np.linspace(0, 7, 100), name=\"y\", dims=\"y\", attrs=dict(units=\"m\"))\n",
    "z = 0\n",
    "# Add coordinates:\n",
    "x.coords[\"x\"] = x\n",
    "y.coords[\"y\"] = y\n",
    "\n",
    "\n",
    "mapped = ds.emc3.evaluate_at_xyz(\n",
    "    x, y, z, [\"ne\", \"Te\"], periodicity=5, updownsym=True, delta_phi=np.pi / 180\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the pressure\n",
    "plt.figure()\n",
    "(mapped.ne * mapped.Te).plot(x=\"x\", y=\"y\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Precalculate the mapping\n",
    "\n",
    "Rather then directly evaluating a quantity, we can also calculate the mapping first. This is usefull if the same mapping is needed for different simulations, as long as they use the identical grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate an abitrary 2D slice\n",
    "x = xr.DataArray(np.linspace(0, 7, 100), name=\"x\", dims=\"x\", attrs=dict(units=\"m\"))\n",
    "y = xr.DataArray(np.linspace(0, 7, 100), name=\"y\", dims=\"y\", attrs=dict(units=\"m\"))\n",
    "z = 0\n",
    "# Add coordinates:\n",
    "x.coords[\"x\"] = x\n",
    "y.coords[\"y\"] = y\n",
    "\n",
    "# We don't pass in the key we want to evaluate\n",
    "# In this case we get a dataset with indices\n",
    "mapped = ds.emc3.evaluate_at_xyz(\n",
    "    x, y, z, key=None, periodicity=5, updownsym=True, delta_phi=np.pi / 180\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "# This time we calculate the pressure\n",
    "ds[\"Pe\"] = ds.Te * ds.ne\n",
    "# And then we evaluate with the pre-evaluated indices the pressure in the plane\n",
    "plt.figure()\n",
    "ds[\"Pe\"].isel(**mapped).plot(x=\"x\", y=\"y\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate along a field line\n",
    "\n",
    "xemc3 doesn't support field-line tracing\n",
    "\n",
    "We can however use the webservices, if we are on the IPP network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "## Warning\n",
    "\n",
    "Please note that if you don't use the same magnetic configuration in the fieldlinetracer and the simulations, you might get slightly wrong results or completely wrong results.\n",
    "    \n",
    "Especially, conclusion drawn might be completely wrong!\n",
    "    \n",
    "<b>Always double check the magnetic configurations!</b>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from dave_utils import jafw\n",
    "except ImportError:\n",
    "    print(\"Importing failed. Consider installing dave_utils with:\")\n",
    "    print(\"pip install --user git+https://gitlab.mpcdf.mpg.de/dave/dave_utils.git\")\n",
    "    jafw = None\n",
    "    raise\n",
    "\n",
    "# This will fail outside of the IPP network\n",
    "flt = jafw.getSrv()\n",
    "\n",
    "config = jafw.setCurrents([-1.74e6] * 5 + [0] * 5)\n",
    "\n",
    "pos = flt.types.Points3D()\n",
    "pos.x1 = 5.7\n",
    "pos.x2 = 0\n",
    "pos.x3 = 0\n",
    "\n",
    "lineTask = flt.types.LineTracing()\n",
    "lineTask.numSteps = 3000\n",
    "\n",
    "task = flt.types.Task()\n",
    "task.step = 0.01\n",
    "task.lines = lineTask\n",
    "\n",
    "res = flt.service.trace(pos, config, task, None, None)\n",
    "\n",
    "i = 0\n",
    "dat = np.array(\n",
    "    [\n",
    "        res.lines[i].vertices.x1,\n",
    "        res.lines[i].vertices.x2,\n",
    "        res.lines[i].vertices.x3,\n",
    "    ]\n",
    ")\n",
    "mapped = ds.emc3.evaluate_at_xyz(\n",
    "    dat[0],\n",
    "    dat[1],\n",
    "    dat[2],\n",
    "    \"ne\",\n",
    "    periodicity=5,\n",
    "    updownsym=True,\n",
    "    delta_phi=np.pi / 180,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if jafw is not None:\n",
    "    mapped[\"ne\"].plot(figsize=(16, 6))\n",
    "else:\n",
    "    print(\"Fieldlinetracer not available ...\")"
   ]
  }
 ],
 "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": 4
}