{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c6ddef50-6fe0-4f0b-8685-a7918331244f",
   "metadata": {},
   "source": [
    "# Spatially varying diffusion coefficients\n",
    "\n",
    "It is possible to create spatially varying diffusion coefficient input files for emc3_eirene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9a511d-1b05-412e-a93e-91d6dcf5575f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import xemc3\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d65f56c-223e-407b-a8f4-bb9cc429c6cc",
   "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": "code",
   "execution_count": null,
   "id": "edfe94a2-5984-41b2-9bea-f41cd12be40b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vol_dims = [x[0] for x in [ds.r.shape, ds.theta.shape, ds.phi.shape]]\n",
    "print(vol_dims)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8e12036-1af1-40f7-aba4-bf1da65fdeed",
   "metadata": {},
   "source": [
    "## Create the desired profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dddcd0ca-7250-4809-a413-de824e45e7d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = np.linspace(0, np.pi / 5, vol_dims[2])[None, :]\n",
    "theta = np.linspace(0, np.pi * 2, vol_dims[1])[:, None]\n",
    "data = np.exp(-(((phi) * 4) ** 2) - (theta - np.pi * 1) ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "191502d4-207e-4006-afe1-453c6b2bd03b",
   "metadata": {},
   "source": [
    "## Plot what you made"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "785eacd1-1b75-4dd6-aa7f-9a79e2dac837",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.pcolormesh(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ce332f-21fe-41d9-b993-3b7d5bf656dd",
   "metadata": {},
   "source": [
    "## Store in emc3_eirene format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79ae010f-9b70-4615-b673-cd83e0251a55",
   "metadata": {},
   "outputs": [],
   "source": [
    "for base in 1, 10:\n",
    "    d2 = 0.1 + (base - 1) * 0.1 * data\n",
    "    k = \"diff\"\n",
    "    ds[k] = ds.Te\n",
    "    ds[k].attrs = {}\n",
    "    ds[k].data[...] = d2 * 1e4\n",
    "    prefix = f\"pol_tor_v1_base={base}_\"\n",
    "    ds.emc3.to_fort(k, f\"{prefix}_diffusion.fort\", kinetic=True)\n",
    "    ds[k].data *= 3\n",
    "    ds.emc3.to_fort(k, f\"{prefix}_heat.fort\", kinetic=True)"
   ]
  }
 ],
 "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}