{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running TARDIS Model Grids\n",
    "\n",
    "This notebook demonstrates the capabilities of the TARDIS grid. The grid facilitates users running large numbers of TARDIS simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tardis.io.atom_data.util import download_atom_data\n",
    "from tardis.io.atom_data import AtomData\n",
    "\n",
    "import tardis.grid as grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We download the atomic data needed to run a TARDIS simulation\n",
    "download_atom_data('kurucz_cd23_chianti_H_He')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating A Grid\n",
    "\n",
    "There are 2 ways of creating a TARDIS grid: directly from a dataframe, or by defining a set of axes over which the grid should be defined. In both cases, a config.yml file is required. **Note that for the dataframe, the column names are in the form of valid keys in a tardis Configuration dictionary. For the axes, the keys must similarly be valid tardis Configuration keys.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a DataFrame\n",
    "df = pd.read_csv('example_grid.txt')\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a tardis grid directly from a dataframe.\n",
    "grid.tardisGrid(configFile='example.yml', gridFrame=df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an axes dictionary\n",
    "axesdict = {'model.structure.velocity.start':np.arange(10000,15000,1000), \n",
    "            'model.abundances.He':np.arange(0,1,0.1),\n",
    "            'model.abundances.H':np.arange(0,1,0.25)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create a tardis grid from an axes dict using the classmethod.\n",
    "grid.tardisGrid.from_axes(configFile='example.yml', axesdict=axesdict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TARDIS Grid Attributes\n",
    "\n",
    "The TARDIS grid only has 2 attributes. It creates a TARDIS Configuration object from the user specified config.yml file and saves this to the `self.config` attribute. The grid also stores the parameters of every grid point in a Dataframe accessed by `self.grid`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tg = grid.tardisGrid(configFile='example.yml', gridFrame=df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The config generated from the user's yml file.\n",
    "tg.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The user specified grid.\n",
    "tg.grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TARDIS Grid Functionality\n",
    "\n",
    "The TARDIS Grid provides a variety of functions for using the grid to generate new Configurations, return new SimulationState objects, or directly run simulations using the parameters specified by the grid. This functionality is particularly useful for running large numbers of simulations and easily works with JobArrays where the row_index is the JobArray index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Easily get a new TARDIS Configuration object \n",
    "# which has the properties of the base self.config\n",
    "# but with properties modified by a row of the grid.\n",
    "new_grid = tg.grid_row_to_config(row_index=0);\n",
    "print(\"tg.config is the original config:\",tg.config.model.abundances)\n",
    "print(\"The new config is modified: \",new_grid.model.abundances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In case a user needs to make more complicated changes\n",
    "# to the base TARDIS model (i.e. using parameters that \n",
    "# are not valid TARDIS Configuration keys), the grid\n",
    "# can return a SimulationState object for a given row.\n",
    "atom_data = AtomData.from_hdf(tg.config.atom_data)\n",
    "simulation_state = tg.grid_row_to_simulation_state(row_index=0, atomic_data=atom_data)\n",
    "simulation_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Users can easily run a full TARDIS simulation\n",
    "# using the grid.\n",
    "sim = tg.run_sim_from_grid(row_index=0)"
   ]
  }
 ],
 "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}