{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Emission Profiles modeled with CHERAB\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import ticker\n",
    "from raysect.optical import Point3D, Spectrum, Vector3D, World\n",
    "\n",
    "from cherab.phix.plasma import import_plasma\n",
    "from cherab.phix.tools.raytransfer import import_phix_rtc\n",
    "from cherab.phix.tools.visualize import show_phix_profile, show_phix_profiles\n",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "world = World()\n",
    "plasma, eq = import_plasma(world)\n",
    "rtc = import_phix_rtc(world, equilibrium=eq)\n",
    "model = [i for i in plasma.models]\n",
    "pprint(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set sample paramaters\n",
    "dr = rtc.material.dr\n",
    "dz = rtc.material.dz\n",
    "nr = rtc.material.grid_shape[0]\n",
    "nz = rtc.material.grid_shape[2]\n",
    "rmin = rtc.material.rmin\n",
    "zmin = rtc.transform[2, 3]\n",
    "rmax = rmin + dr * nr\n",
    "zmax = zmin + dz * nz\n",
    "\n",
    "xrange = np.linspace(rmin, rmax, nr)\n",
    "yrange = np.linspace(zmin, zmax, nz)\n",
    "direction = Vector3D(0, 1, 0)\n",
    "# H-balmar emission line\n",
    "Halpha = (655.6, 656.8)\n",
    "Hbeta = (485.6, 486.5)\n",
    "Hgamma = (433.6, 434.4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bremsstrahlung emission\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate spectral radiance at $r-z$ plane.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate spectral radiance\n",
    "spectrum_bins = 500\n",
    "min_wavelength = 375  # [nm]\n",
    "max_wavelength = 780\n",
    "radiance = np.zeros((nr, nz))\n",
    "spectral_radiance = np.zeros((spectrum_bins, nr, nz))\n",
    "\n",
    "for i, x in enumerate(xrange):\n",
    "    for j, y in enumerate(yrange):\n",
    "        emission = model[0].emission(\n",
    "            Point3D(x, 0.0, y), direction, Spectrum(min_wavelength, max_wavelength, spectrum_bins)\n",
    "        )\n",
    "        spectral_radiance[:, i, j] = emission.samples\n",
    "        radiance[i, j] = emission.total()\n",
    "\n",
    "# mask values under 0\n",
    "radiance = np.ma.masked_where(radiance <= 0, radiance)\n",
    "spectral_radiance = np.ma.masked_where(spectral_radiance <= 0, spectral_radiance)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "show spacial profile in $r-z$ plane and spectrum at $r=0.33$ line.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, constrained_layout=True, figsize=(10, 4))\n",
    "\n",
    "# show profile\n",
    "show_phix_profile(axes[0], radiance, cmap=\"Reds\", rtc=rtc, plot_contour=False)\n",
    "axes[0].set_title(\"Profile in $r-z$ plane\", pad=15)\n",
    "fmt = ticker.ScalarFormatter(useMathText=True)\n",
    "fmt.set_powerlimits((0, 0))\n",
    "cbar0 = plt.colorbar(axes[0].collections[0], ax=axes[0], format=fmt, pad=0)\n",
    "cbar0.set_label(\"radiance per meter [W/m$^3$/str]\")\n",
    "cbar0.ax.yaxis.set_offset_position(\"left\")\n",
    "axes[0].set_xlabel(\"$R$ [m]\")\n",
    "axes[0].set_ylabel(\"$Z$ [m]\")\n",
    "\n",
    "# show spectrum\n",
    "r0_index = 46  # index of magnetix axis R in xrange\n",
    "pc = axes[1].pcolormesh(\n",
    "    emission.wavelengths,\n",
    "    yrange,\n",
    "    spectral_radiance[:, r0_index, :].T,\n",
    "    cmap=\"Reds\",\n",
    ")\n",
    "fmt = ticker.ScalarFormatter(useMathText=True)\n",
    "fmt.set_powerlimits((0, 0))\n",
    "cbar1 = plt.colorbar(\n",
    "    pc, label=\"Spectral radiance per meter [W/m$^3$/str/nm]\", ax=axes[1], format=fmt, pad=0\n",
    ")\n",
    "cbar1.ax.yaxis.set_offset_position(\"left\")\n",
    "axes[1].set_xlabel(\"wavelength [nm]\")\n",
    "axes[1].set_ylabel(\"$Z $[m]\")\n",
    "axes[1].set_title(f\"Spectrum along the line $r={xrange[r0_index]:.2f}$ [m]\", pad=15)\n",
    "\n",
    "fig.suptitle(\"Bremsstrahlung emission\", x=0.55)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot 1-D bremsstrahlung spectrum at $(r, z) = (0.34, 0)$ which is magnetic axis position.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# spectral_radiance index at (r, z) = magnetic axis\n",
    "idx = np.abs(xrange - eq.magnetic_axis.x).argmin()\n",
    "idy = np.where(yrange == 0)[0][0]\n",
    "\n",
    "# plot 1-d spectrum\n",
    "fig, ax = plt.subplots(1, constrained_layout=True)\n",
    "ax.plot(emission.wavelengths, spectral_radiance[:, idx, idy])\n",
    "ax.set_xlabel(\"wavelength [nm]\")\n",
    "ax.set_ylabel(\"Spectral radiance per meter [W/m$^3$/str/nm]\")\n",
    "ax.yaxis.set_major_formatter(ticker.ScalarFormatter(useMathText=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Excitation emission of hydrogen balmar-series\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum_bins = 50\n",
    "radiance = np.zeros((3, nr, nz))\n",
    "spectral_radiance = np.zeros((3, spectrum_bins, nr, nz))  # number of Hydrogen lines (ax=0)\n",
    "for i, x in enumerate(xrange):\n",
    "    for j, y in enumerate(yrange):\n",
    "        for k, waverange in enumerate([Halpha, Hbeta, Hgamma]):\n",
    "            emission = model[k + 1].emission(\n",
    "                Point3D(x, 0.0, y), direction, Spectrum(waverange[0], waverange[1], spectrum_bins)\n",
    "            )\n",
    "            spectral_radiance[k, :, i, j] = emission.samples\n",
    "            radiance[k, i, j] = emission.total()\n",
    "\n",
    "# mask radiance values under 0\n",
    "radiance = np.ma.masked_where(radiance <= 0, radiance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, grids = show_phix_profiles(\n",
    "    [radiance[k, :] for k in range(len(radiance[:, 0, 0]))],\n",
    "    clabel=\"Radiance per meter [W/m$^3$/str]\",\n",
    "    cmap=\"Reds\",\n",
    "    rtc=rtc,\n",
    "    cbar_mode=\"single\",\n",
    "    plot_mode=\"scalar\",\n",
    "    axes_pad=0.02,\n",
    ")\n",
    "for grid, name in zip(grids, [\"$_\\\\alpha$\", \"$_\\\\beta$\", \"$_\\\\gamma$\"]):\n",
    "    grid.set_title(f\"H{name}\")\n",
    "fig.suptitle(\"Excitation emission of hydrogen balmer-series\", y=0.93)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, grids = show_phix_profiles(\n",
    "    [radiance[k, :] for k in range(len(radiance[:, 0, 0]))],\n",
    "    clabel=\"Radiance per meter [W/m$^3$/str]\",\n",
    "    cmap=\"Reds\",\n",
    "    rtc=rtc,\n",
    "    cbar_mode=\"each\",\n",
    "    plot_mode=\"eng\",\n",
    "    axes_pad=0.5,\n",
    ")\n",
    "for grid, name in zip(grids, [\"$_\\\\alpha$\", \"$_\\\\beta$\", \"$_\\\\gamma$\"]):\n",
    "    grid.set_title(f\"H{name}\")\n",
    "fig.suptitle(\"Excitation emission of hydrogen balmer-series\", y=0.86)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recombination emission of hydrogen balmar-series\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum_bins = 50\n",
    "radiance = np.zeros((3, nr, nz))\n",
    "spectral_radiance = np.zeros((3, spectrum_bins, nr, nz))  # number of Hydrogen lines (ax=0)\n",
    "for i, x in enumerate(xrange):\n",
    "    for j, y in enumerate(yrange):\n",
    "        for k, waverange in enumerate([Halpha, Hbeta, Hgamma]):\n",
    "            emission = model[k + 5].emission(\n",
    "                Point3D(x, 0.0, y), direction, Spectrum(waverange[0], waverange[1], spectrum_bins)\n",
    "            )\n",
    "            spectral_radiance[k, :, i, j] = emission.samples\n",
    "            radiance[k, i, j] = emission.total()\n",
    "\n",
    "# mask radiance values under 0\n",
    "radiance = np.ma.masked_where(radiance <= 0, radiance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, grids = show_phix_profiles(\n",
    "    [radiance[k, :] for k in range(3)],\n",
    "    clabel=\"Radiance per meter [W/m$^3$/str]\",\n",
    "    cmap=\"Reds\",\n",
    "    rtc=rtc,\n",
    "    cbar_mode=\"single\",\n",
    ")\n",
    "for grid, name in zip(grids, [\"$_\\\\alpha$\", \"$_\\\\beta$\", \"$_\\\\gamma$\"]):\n",
    "    grid.set_title(f\"H{name}\")\n",
    "fig.suptitle(\"Recombination emission of hydrogen balmer-series\", y=0.93)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, grids = show_phix_profiles(\n",
    "    [radiance[k, :] for k in range(3)],\n",
    "    clabel=\"Radiance per meter [W/m$^3$/str]\",\n",
    "    cmap=\"Reds\",\n",
    "    rtc=rtc,\n",
    "    cbar_mode=\"each\",\n",
    "    axes_pad=0.5,\n",
    ")\n",
    "for grid, name in zip(grids, [\"$_\\\\alpha$\", \"$_\\\\beta$\", \"$_\\\\gamma$\"]):\n",
    "    grid.set_title(f\"H{name}\")\n",
    "fig.suptitle(\"Recombination emission of hydrogen balmer-series\", y=0.86)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Spectral Radiance per meter at magnetic axis in hydrogen balmar-series\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum_bins = 50\n",
    "spectral_radiance_e = np.zeros((3, spectrum_bins))\n",
    "spectral_radiance_r = np.zeros((3, spectrum_bins))\n",
    "for k, waverange in enumerate([Halpha, Hbeta, Hgamma]):\n",
    "    emission_e = model[k + 1].emission(\n",
    "        Point3D(eq.magnetic_axis.x, 0.0, eq.magnetic_axis.y),\n",
    "        direction,\n",
    "        Spectrum(waverange[0], waverange[1], spectrum_bins),\n",
    "    )\n",
    "    emission_r = model[k + 5].emission(\n",
    "        Point3D(eq.magnetic_axis.x, 0.0, eq.magnetic_axis.y),\n",
    "        direction,\n",
    "        Spectrum(waverange[0], waverange[1], spectrum_bins),\n",
    "    )\n",
    "    spectral_radiance_e[k, :] = emission_e.samples\n",
    "    spectral_radiance_r[k, :] = emission_r.samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength = np.concatenate(\n",
    "    [\n",
    "        np.linspace(*Hgamma, len(spectral_radiance[2])),\n",
    "        np.linspace(*Hbeta, len(spectral_radiance[1])),\n",
    "        np.linspace(*Halpha, len(spectral_radiance[0])),\n",
    "    ]\n",
    ")\n",
    "spectrum_ex = np.concatenate([spectral_radiance_e[i, :] for i in [2, 1, 0]])\n",
    "spectrum_re = np.concatenate([spectral_radiance_r[i, :] for i in [2, 1, 0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(wavelength, spectrum_ex, label=\"excitation\")\n",
    "# ax.plot(wavelength, spectrum_re, label=\"recombination\")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"wavelength [nm]\")\n",
    "ax.set_ylabel(\"Spectral radiance per meter [W/m$^3$/str/nm]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "# ax.plot(wavelength, spectrum_ex, label=\"excitation\")\n",
    "ax.plot(wavelength, spectrum_re, label=\"recombination\", color=\"C1\")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"wavelength [nm]\")\n",
    "ax.set_ylabel(\"Spectral radiance per meter [W/m$^3$/str/nm]\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "cherab-phix-dev",
   "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.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "2725905a4c02db19e04df9b8fdbbe5ec65a73ea52bebaf9474aa1cc98819834c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}