{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EasyVVUQ fusion tutorial\n",
    "\n",
    "Run an EasyVVUQ campaign to analyze the sensitivity of the temperature\n",
    "profile predicted by a simplified model of heat conduction in a\n",
    "tokamak plasma.\n",
    "\n",
    "This is done with PCE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install EasyVVUQ\n",
    "!pip install future\n",
    "!pip install fipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages that we will use\n",
    "\n",
    "import os\n",
    "import easyvvuq as uq\n",
    "import chaospy as cp\n",
    "import pickle\n",
    "import time\n",
    "import numpy as np \n",
    "import matplotlib.pylab as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a fresh campaign called \"fusion_pce.\"\n",
    "\n",
    "my_campaign = uq.Campaign(name='fusion_pce.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define parameter space\n",
    "\n",
    "params = {\n",
    "    \"Qe_tot\":   {\"type\": \"float\",   \"min\": 1.0e6, \"max\": 50.0e6, \"default\": 2e6}, \n",
    "    \"H0\":       {\"type\": \"float\",   \"min\": 0.00,  \"max\": 1.0,    \"default\": 0}, \n",
    "    \"Hw\":       {\"type\": \"float\",   \"min\": 0.01,  \"max\": 100.0,  \"default\": 0.1}, \n",
    "    \"Te_bc\":    {\"type\": \"float\",   \"min\": 10.0,  \"max\": 1000.0, \"default\": 100}, \n",
    "    \"chi\":      {\"type\": \"float\",   \"min\": 0.01,  \"max\": 100.0,  \"default\": 1}, \n",
    "    \"a0\":       {\"type\": \"float\",   \"min\": 0.2,   \"max\": 10.0,   \"default\": 1}, \n",
    "    \"R0\":       {\"type\": \"float\",   \"min\": 0.5,   \"max\": 20.0,   \"default\": 3}, \n",
    "    \"E0\":       {\"type\": \"float\",   \"min\": 1.0,   \"max\": 10.0,   \"default\": 1.5}, \n",
    "    \"b_pos\":    {\"type\": \"float\",   \"min\": 0.95,  \"max\": 0.99,   \"default\": 0.98}, \n",
    "    \"b_height\": {\"type\": \"float\",   \"min\": 3e19,  \"max\": 10e19,  \"default\": 6e19}, \n",
    "    \"b_sol\":    {\"type\": \"float\",   \"min\": 2e18,  \"max\": 3e19,   \"default\": 2e19}, \n",
    "    \"b_width\":  {\"type\": \"float\",   \"min\": 0.005, \"max\": 0.025,  \"default\": 0.01}, \n",
    "    \"b_slope\":  {\"type\": \"float\",   \"min\": 0.0,   \"max\": 0.05,   \"default\": 0.01}, \n",
    "    \"nr\":       {\"type\": \"integer\", \"min\": 10,    \"max\": 1000,   \"default\": 100}, \n",
    "    \"dt\":       {\"type\": \"float\",   \"min\": 1e-3,  \"max\": 1e3,    \"default\": 100},\n",
    "    \"out_file\": {\"type\": \"string\",  \"default\": \"output.csv\"}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an encoder, decoder and collater for PCE test app\n",
    "\n",
    "encoder = uq.encoders.GenericEncoder(template_fname='fusion.template',\n",
    "                                     delimiter='$',\n",
    "                                     target_filename='fusion_in.json')\n",
    "\n",
    "\n",
    "decoder = uq.decoders.SimpleCSV(target_filename=\"output.csv\",\n",
    "                                output_columns=[\"te\", \"ne\", \"rho\", \"rho_norm\"])\n",
    "                                \n",
    "\n",
    "#collater = uq.collate.AggregateSamples(average=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the app (automatically set as current app)\n",
    "\n",
    "my_campaign.add_app(name=\"fusion\",\n",
    "                    params=params,\n",
    "                    encoder=encoder,\n",
    "                    decoder=decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the sampler (here simplified to two uncertain quantities)\n",
    "\n",
    "vary = {\n",
    "    \"Qe_tot\":   cp.Uniform(1.8e6, 2.2e6),\n",
    "    \"Te_bc\":    cp.Uniform(80.0,  120.0)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"\"\" other possible quantities to vary\n",
    "    \"H0\":       cp.Uniform(0.0,   0.2),\n",
    "    \"Hw\":       cp.Uniform(0.1,   0.5),\n",
    "    \"chi\":      cp.Uniform(0.8,   1.2), \n",
    "\n",
    "    \"a0\":       cp.Uniform(0.9,   1.1), \n",
    "    \"R0\":       cp.Uniform(2.7,   3.3), \n",
    "    \"E0\":       cp.Uniform(1.4,   1.6), \n",
    "    \"b_pos\":    cp.Uniform(0.95,  0.99), \n",
    "    \"b_height\": cp.Uniform(5e19,  7e19), \n",
    "    \"b_sol\":    cp.Uniform(1e19,  3e19), \n",
    "    \"b_width\":  cp.Uniform(0.015, 0.025), \n",
    "    \"b_slope\":  cp.Uniform(0.005, 0.020)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Associate a sampler with the campaign\n",
    "\n",
    "my_campaign.set_sampler(uq.sampling.PCESampler(vary=vary, polynomial_order=3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Will draw all (of the finite set of samples)\n",
    "\n",
    "my_campaign.draw_samples()\n",
    "print('Number of samples = %s' % my_campaign.get_active_sampler().count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and populate the run directories\n",
    "\n",
    "my_campaign.populate_runs_dir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the cases\n",
    "\n",
    "cwd = os.getcwd().replace(' ', '\\ ')      # deal with ' ' in the path\n",
    "cmd = f\"{cwd}/fusion_model.py fusion_in.json\"\n",
    "print(cmd)\n",
    "my_campaign.apply_for_each_run_dir(uq.actions.ExecuteLocal(cmd, interpret='python3'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collate the results\n",
    "\n",
    "my_campaign.collate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Post-processing analysis\n",
    "\n",
    "my_campaign.apply_analysis(uq.analysis.PCEAnalysis(sampler=my_campaign.get_active_sampler(), \n",
    "                                                   qoi_cols=[\"te\", \"ne\", \"rho\", \"rho_norm\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get Descriptive Statistics\n",
    "\n",
    "results = my_campaign.get_last_analysis()\n",
    "rho = results.describe('rho', 'mean')\n",
    "rho_norm = results.describe('rho_norm', 'mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the calculated Te: mean, with std deviation, 10 and 90% and range\n",
    "plt.figure() \n",
    "plt.plot(rho, results.describe('te', 'mean'), 'b-', label='Mean')\n",
    "plt.plot(rho, results.describe('te', 'mean')-results.describe('te', 'std'), 'b--', label='+1 std deviation')\n",
    "plt.plot(rho, results.describe('te', 'mean')+results.describe('te', 'std'), 'b--')\n",
    "plt.fill_between(rho, results.describe('te', 'mean')-results.describe('te', 'std'), results.describe('te', 'mean')+results.describe('te', 'std'), color='b', alpha=0.2)\n",
    "plt.plot(rho, results.describe('te', '10%'), 'b:', label='10 and 90 percentiles')\n",
    "plt.plot(rho, results.describe('te', '90%'), 'b:')\n",
    "plt.fill_between(rho, results.describe('te', '10%'), results.describe('te', '90%'), color='b', alpha=0.1)\n",
    "plt.fill_between(rho, results.describe('te', 'min'), results.describe('te', 'max'), color='b', alpha=0.05)\n",
    "plt.legend(loc=0)\n",
    "plt.xlabel('rho [m]')\n",
    "plt.ylabel('Te [eV]')\n",
    "plt.title(my_campaign.campaign_dir);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the first Sobol results\n",
    "\n",
    "plt.figure() \n",
    "for k in results.sobols_first()['te'].keys(): plt.plot(rho, results.sobols_first()['te'][k], label=k)\n",
    "plt.legend(loc=0)\n",
    "plt.xlabel('rho [m]')\n",
    "plt.ylabel('sobols_first')\n",
    "plt.title(my_campaign.campaign_dir);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the second Sobol results\n",
    "\n",
    "plt.figure() \n",
    "for k1 in results.sobols_second()['te'].keys(): \n",
    "    for k2 in results.sobols_second()['te'][k1].keys():\n",
    "        plt.plot(rho, results.sobols_second()['te'][k1][k2], label=k1+'/'+k2)\n",
    "plt.legend(loc=0)    \n",
    "plt.xlabel('rho [m]')\n",
    "plt.ylabel('sobols_second')\n",
    "plt.title(my_campaign.campaign_dir+'\\n');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the total Sobol results\n",
    "\n",
    "plt.figure() \n",
    "for k in results.sobols_total()['te'].keys(): plt.plot(rho, results.sobols_total()['te'][k], label=k)\n",
    "plt.legend(loc=0)    \n",
    "plt.xlabel('rho [m]')\n",
    "plt.ylabel('sobols_total')\n",
    "plt.title(my_campaign.campaign_dir);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the distributions\n",
    "\n",
    "plt.figure() \n",
    "for i, D in enumerate(results.raw_data['output_distributions']['te']):\n",
    "    _Te = np.linspace(D.lower[0], D.upper[0], 101)\n",
    "    _DF = D.pdf(_Te)\n",
    "    plt.loglog(_Te, _DF, 'b-', alpha=0.25)\n",
    "    plt.loglog(results.describe('te', 'mean')[i], np.interp(results.describe('te', 'mean')[i], _Te, _DF), 'bo')\n",
    "    plt.loglog(results.describe('te', 'mean')[i]-results.describe('te', 'std')[i], np.interp(results.describe('te', 'mean')[i]-results.describe('te', 'std')[i], _Te, _DF), 'b*')\n",
    "    plt.loglog(results.describe('te', 'mean')[i]+results.describe('te', 'std')[i], np.interp(results.describe('te', 'mean')[i]+results.describe('te', 'std')[i], _Te, _DF), 'b*')\n",
    "    plt.loglog(results.describe('te', '10%')[i],  np.interp(results.describe('te', '10%')[i], _Te, _DF), 'b+')\n",
    "    plt.loglog(results.describe('te', '90%')[i],  np.interp(results.describe('te', '90%')[i], _Te, _DF), 'b+')\n",
    "plt.xlabel('Te')\n",
    "plt.ylabel('distribution function');"
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}