TARDIS Grid (tardis-sn#1738)

* experimental notebook * init grid dataframe * config changes * added grid module * updated test * added docstrings * Added grid_row_to_model() class function. * Added docs for grid [build docs] * [build docs] * [build docs] * [build docs] * [build docs] * added grid doc to sidebar. [build docs] * Added tests * black formatting * renamed to base.py * cleared output so notebook is run when building docs * added kwarg pass through for run_tardis kwargs. * Fixed docstring Co-authored-by: Marc Williamson <[email protected]>
atharva-2001 · Jul 28, 2021 · a6522b0 · a6522b0
1 parent d7d172b
commit a6522b0
Show file tree

Hide file tree

Showing 10 changed files with 587 additions and 1 deletion.
diff --git a/docs/index.rst b/docs/index.rst
@@ -62,7 +62,7 @@ Mission Statement
     io/optional/index
     io/visualization/index
     io/output/index
-
+    io/grid/TardisGridTutorial
 
 .. toctree::
     :maxdepth: 2

diff --git a/docs/io/grid/TardisGridTutorial.ipynb b/docs/io/grid/TardisGridTutorial.ipynb
@@ -0,0 +1,184 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# TARDIS Grid Tutorial\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",
+    "\n",
+    "import tardis.grid as grid"
+   ]
+  },
+  {
+   "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 Radial1DModel 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 Radial1DModel object for a given row.\n",
+    "model = tg.grid_row_to_model(row_index=0)\n",
+    "model"
+   ]
+  },
+  {
+   "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
+}
diff --git a/docs/io/grid/example.yml b/docs/io/grid/example.yml
@@ -0,0 +1,57 @@
+# Example YAML configuration for TARDIS
+tardis_config_version: v1.0
+
+supernova:
+  luminosity_requested: 5.679e41 erg/s
+  luminosity_wavelength_start: 3481.82000178 angstrom
+  luminosity_wavelength_end: 9947.78776065 angstrom
+  time_explosion: 10 day
+
+atom_data: kurucz_cd23_chianti_H_He.h5
+
+model:
+  structure:
+    type: specific
+    velocity:
+      start: 1.0e4 km/s
+      stop: 20000 km/s
+      num: 20
+    density:
+      type: branch85_w7
+
+  abundances:
+    type: uniform
+    He: 0.2
+    H: 0.8
+
+
+plasma:
+  initial_t_inner: 7000 K
+  disable_electron_scattering: no
+  ionization: nebular
+  excitation: dilute-lte
+  radiative_rates_type: dilute-blackbody
+  line_interaction_type: macroatom
+
+montecarlo:
+  seed: 23111963
+  no_of_packets: 4.0e+4
+  iterations: 2
+  nthreads: 1
+
+  last_no_of_packets: 1.e+3
+  no_of_virtual_packets: 3
+
+  convergence_strategy:
+    type: damped
+    damping_constant: 0.5
+    threshold: 0.05
+    fraction: 0.8
+    hold_iterations: 3
+    t_inner:
+      damping_constant: 0.5
+
+spectrum:
+  start: 500 angstrom
+  stop: 20000 angstrom
+  num: 10000
diff --git a/docs/io/grid/example_grid.txt b/docs/io/grid/example_grid.txt
@@ -0,0 +1,4 @@
+model.abundances.H,model.abundances.He,model.structure.velocity.start,plasma.initial_t_inner
+1.0,0.0,10000,5000
+0.4,0.6,12000,6000
+0.7,0.3,15000,7000
diff --git a/tardis/grid/__init__.py b/tardis/grid/__init__.py
@@ -0,0 +1 @@
+from tardis.grid.base import *