From a6522b045582b9baf1ea0c2f80f216b623678899 Mon Sep 17 00:00:00 2001 From: Marc Williamson Date: Wed, 28 Jul 2021 14:01:41 -0400 Subject: [PATCH] TARDIS Grid (#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 --- docs/index.rst | 2 +- docs/io/grid/TardisGridTutorial.ipynb | 184 ++++++++++++++++++++++++ docs/io/grid/example.yml | 57 ++++++++ docs/io/grid/example_grid.txt | 4 + tardis/grid/__init__.py | 1 + tardis/grid/base.py | 172 ++++++++++++++++++++++ tardis/grid/test.yml | 65 +++++++++ tardis/grid/tests/data/example.yml | 57 ++++++++ tardis/grid/tests/data/example_grid.txt | 4 + tardis/grid/tests/test_grid.py | 42 ++++++ 10 files changed, 587 insertions(+), 1 deletion(-) create mode 100644 docs/io/grid/TardisGridTutorial.ipynb create mode 100644 docs/io/grid/example.yml create mode 100644 docs/io/grid/example_grid.txt create mode 100644 tardis/grid/__init__.py create mode 100644 tardis/grid/base.py create mode 100644 tardis/grid/test.yml create mode 100644 tardis/grid/tests/data/example.yml create mode 100644 tardis/grid/tests/data/example_grid.txt create mode 100644 tardis/grid/tests/test_grid.py diff --git a/docs/index.rst b/docs/index.rst index 9e36dd67740..360ff94c108 100644 --- 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 new file mode 100644 index 00000000000..bc4393bb4b6 --- /dev/null +++ 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 new file mode 100644 index 00000000000..2161d31f9bf --- /dev/null +++ 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 new file mode 100644 index 00000000000..86815cf0e59 --- /dev/null +++ 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 new file mode 100644 index 00000000000..eb8bb8410b2 --- /dev/null +++ b/tardis/grid/__init__.py @@ -0,0 +1 @@ +from tardis.grid.base import * diff --git a/tardis/grid/base.py b/tardis/grid/base.py new file mode 100644 index 00000000000..28ac97bf864 --- /dev/null +++ b/tardis/grid/base.py @@ -0,0 +1,172 @@ +import pandas as pd +import numpy as np +import copy +import tardis + +from tardis.io.config_reader import Configuration +from tardis.model import Radial1DModel + + +def _set_tardis_config_property(tardis_config, key, value): + """ + Sets tardis_config[key] = value for a TARDIS config + dictionary. Recursively searches for the deepest + dictionary and sets the corresponding value. + + Parameters + ---------- + tardis_config : tardis.io.config_reader.Configuration + The config where the key,value pair should be set. + key : str + The key in the tardis_config. Should be of the format + 'property.property.property' + value : object + The value to assign to the config dict for the + corresponding key. + """ + keyitems = key.split(".") + tmp_dict = getattr(tardis_config, keyitems[0]) + for key in keyitems[1:-1]: + tmp_dict = getattr(tmp_dict, key) + setattr(tmp_dict, keyitems[-1], value) + return + + +class tardisGrid: + """ + A class that stores a grid of TARDIS parameters and + facilitates running large numbers of simulations + easily. + + Parameters + ---------- + configFile : str or dict + path to TARDIS yml file, or a pre-validated config dictionary. + gridFrame : pandas.core.frame.DataFrame + dataframe where each row is a set of parameters for + a TARDIS simulation. + + Attributes + ---------- + config : tardis.io.config_reader.Configuration + The validated config dict read from the user provided + configFile. This provides the base properties for the + TARDIS simulation which the rows of the grid modify. + grid : pandas.core.frame.DataFrame + Dataframe where each row is a set of parameters for + a TARDIS simulation. + """ + + def __init__(self, configFile, gridFrame): + try: + tardis_config = Configuration.from_yaml(configFile) + except TypeError: + tardis_config = Configuration.from_config_dict(configFile) + + self.config = tardis_config + self.grid = gridFrame + + return + + def grid_row_to_config(self, row_index): + """ + Converts a grid row to a TARDIS config dict. + Modifies the base self.config according to the + row in self.grid accessed at the provided row_index. + Returns a deep copy so that the base config is + not changed. + + Parameters + ---------- + row_index : int + Row index in grid. + + Returns + ------- + tmp_config : tardis.io.config_reader.Configuration + Deep copy of the base self.config with modified + properties according to the selected row in the grid. + """ + tmp_config = copy.deepcopy(self.config) + grid_row = self.grid.iloc[row_index] + for colname, value in zip(self.grid.columns, grid_row.values): + _set_tardis_config_property(tmp_config, colname, value) + return tmp_config + + def grid_row_to_model(self, row_index): + """ + Generates a TARDIS Radial1DModel object using the base + self.config modified by the specified grid row. + + Parameters + ---------- + row_index : int + Row index in grid. + + Returns + ------- + model : tardis.model.base.Radial1DModel + """ + rowconfig = self.grid_row_to_config(row_index) + model = Radial1DModel.from_config(rowconfig) + return model + + def run_sim_from_grid(self, row_index, **tardiskwargs): + """ + Runs a full TARDIS simulation using the base self.config + modified by the user specified row_index. + + Parameters + ---------- + row_index : int + Row index in grid. + + Returns + ------- + sim : tardis.simulation.base.Simulation + Completed TARDIS simulation object. + """ + tardis_config = self.grid_row_to_config(row_index) + sim = tardis.run_tardis(tardis_config, **tardiskwargs) + return sim + + def save_grid(self, filename): + """ + Saves the parameter grid. Does not save the base + self.config in any way. + + Parameters + ---------- + filename : str + File name to save grid. + """ + self.grid.to_csv(filename, index=False) + return + + @classmethod + def from_axes(cls, configFile, axesdict): + """ + Creates a grid from a set of axes. The axes are provided + as a dictionary, where each key is a valid tardis config + key, and the value is an iterable of values. + + Parameters + ---------- + configFile : str + Path to TARDIS yml file. + axesdict : dict() + Dictionary containing tardis config keys and the + corresponding values to define a grid of tardis + parameters. + """ + axes = [] + dim = 1 + for key in axesdict: + ax = axesdict[key] + axes.append(ax) + dim = dim * len(ax) + axesmesh = np.meshgrid(*axes) + tmp = np.dstack(axesmesh) + gridpoints = tmp.reshape((dim, len(axes)), order="F") + df = pd.DataFrame(data=gridpoints, columns=axesdict.keys()) + return cls(configFile=configFile, gridFrame=df) diff --git a/tardis/grid/test.yml b/tardis/grid/test.yml new file mode 100644 index 00000000000..89a627c3e0d --- /dev/null +++ b/tardis/grid/test.yml @@ -0,0 +1,65 @@ +# 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: /Users/marcwilliamson/Research/TARDIS/tardis-refdata/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.6658 + H: 0.0 + N: 0.05 + Fe: 0.015 + O: 0.00 + Ni: 0.2505 + C: 0.008 + Si: 0.001 + Ca: 0.0102 + + +plasma: + initial_t_inner: 7000 K + disable_electron_scattering: no + ionization: nebular + excitation: dilute-lte + radiative_rates_type: dilute-blackbody + line_interaction_type: macroatom + helium_treatment: recomb-nlte + +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/tardis/grid/tests/data/example.yml b/tardis/grid/tests/data/example.yml new file mode 100644 index 00000000000..2161d31f9bf --- /dev/null +++ b/tardis/grid/tests/data/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/tardis/grid/tests/data/example_grid.txt b/tardis/grid/tests/data/example_grid.txt new file mode 100644 index 00000000000..86815cf0e59 --- /dev/null +++ b/tardis/grid/tests/data/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/tests/test_grid.py b/tardis/grid/tests/test_grid.py new file mode 100644 index 00000000000..5676c6782fb --- /dev/null +++ b/tardis/grid/tests/test_grid.py @@ -0,0 +1,42 @@ +import pandas as pd +import numpy as np +import tardis +import os +import pytest +import tardis.grid as grid + + +DATA_PATH = os.path.join(tardis.__path__[0], "grid", "tests", "data") + + +def test_grid(): + """Tests the basic functionality of the TARDIS grid module.""" + dfpath = os.path.join(DATA_PATH, "example_grid.txt") + ymlpath = os.path.join(DATA_PATH, "example.yml") + axesdict = { + "model.structure.velocity.start": np.arange(10000, 15000, 1000), + "model.abundances.He": np.arange(0, 1, 0.1), + "model.abundances.H": np.arange(0, 1, 0.25), + } + + df = pd.read_csv(os.path.join(dfpath)) + g = grid.tardisGrid(configFile=ymlpath, gridFrame=df) + g2 = grid.tardisGrid.from_axes(configFile=ymlpath, axesdict=axesdict) + + # Check that grid attribute has the right shape + assert g.grid.shape == df.shape + ax_len = 1 + for key in axesdict: + ax_len *= len(axesdict[key]) + assert g2.grid.shape[0] == ax_len + + newconf = g.grid_row_to_config(row_index=0) + # Verify that grid_row_to_config modifies the base config attribute + assert g.config != newconf + + # Verify that a model can be returned. + model = g.grid_row_to_model(row_index=0) + assert ( + model.velocity[0].to("km/s").value + == df.iloc[0]["model.structure.velocity.start"] + )