Add capability to store information about the plasma state during ite…

…rations (#848) * Add logging of W during iterations * Add t_rad and t_inner to logging * Fix typo * Add reshaping of plasma state storage * Move new features to mixin class * Add information about iterations to HDF storage * Add tests for plasma state iterations * Add explicit unit test for reshaping iteration stores * Fix shape typo * Add docstrings * Add unit test for store_plasma_state * Change to new reference data in TRAVIS * Add and update documentation * Fix typo * Switch back to ref-data master branch in Travis * Fix typo
tardis-sn · Aug 3, 2018 · b432f8a · b432f8a
1 parent e91a7b7
commit b432f8a
Show file tree

Hide file tree

Showing 12 changed files with 207 additions and 7 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -74,8 +74,8 @@ before_install:
     - if [[ $TEST_MODE == 'spectrum' ]]; then git lfs install --skip-smudge; fi
     - if [[ $TEST_MODE == 'spectrum' ]]; then git clone $TARDIS_REF_DATA_URL $HOME/tardis-refdata; fi
     - if [[ $TEST_MODE == 'spectrum' ]]; then cd $HOME/tardis-refdata; fi
-    - if [[ $TEST_MODE == 'spectrum' ]]; then git fetch origin pull/3/head:carsus-ref; fi
-    - if [[ $TEST_MODE == 'spectrum' ]]; then git checkout carsus-ref; fi
+    - if [[ $TEST_MODE == 'spectrum' ]]; then git fetch origin; fi
+    - if [[ $TEST_MODE == 'spectrum' ]]; then git checkout origin/master; fi
     - if [[ $TEST_MODE == 'spectrum' ]]; then git lfs pull --include="atom_data/kurucz_cd23_chianti_H_He.h5" origin; fi
     - if [[ $TEST_MODE == 'spectrum' ]]; then git lfs pull --include="atom_data/chianti_He.h5" origin; fi
     - if [[ $TEST_MODE == 'spectrum' ]]; then git lfs pull --include="plasma_reference/" origin; fi

diff --git a/docs/configuration/index.rst b/docs/configuration/index.rst
@@ -15,4 +15,4 @@ and pages provide more details concerning the TARDIS configuration process.
 
   configuration
   config_validator
-  read_configuration.ipynb
+  read_configuration.ipynb
diff --git a/docs/index.rst b/docs/index.rst
@@ -37,8 +37,8 @@ Using Tardis
     :maxdepth: 2
 
     installation
-    running
-    configuration/index
+    quickstart
+    running/index
     examples/index
     scripts/index
     credits

diff --git a/docs/running.rst → docs/quickstart.rst b/docs/running.rst → docs/quickstart.rst
diff --git a/docs/running/access.rst b/docs/running/access.rst
@@ -0,0 +1,15 @@
+******************************
+Accessing Physical Information
+******************************
+
+Various information about numerical and physical properties of a Tardis
+simulation are stored during a run and can be accessed afterwards. Here, we
+describe how some of this information, which we deem most useful for typical
+use cases, can be accessed.
+
+.. toctree::
+
+  access_spectra
+  access_final_plasma
+  access_model
+  access_iterations
diff --git a/docs/running/access_iterations.rst b/docs/running/access_iterations.rst
@@ -0,0 +1,24 @@
+**********************************************
+Access information about individual iterations
+**********************************************
+
+Currently we store information about the plasma state and the inner boundary
+for each iteration. This is saved in the simulation object and can be accessed
+via
+
+.. code-block:: python
+
+  import tardis
+  mdl = tardis.run_tardis("tardis_config.yml")
+  mdl.iterations_w
+
+in case of the dilution factor. 
+
+Currently, the following properties are available:
+
+.. code-block:: python
+
+   mdl.iterations_w  # dilution factor in each cell
+   mdl.iterations_t_rads  # radiation temperature in each cell
+   mdl.iterations_electron_densities  # electron density in each cell
+   mdl.iterations_t_inner  # inner boundary temperature
diff --git a/docs/running/access_model.rst b/docs/running/access_model.rst
diff --git a/docs/running/access_plasma.rst b/docs/running/access_plasma.rst
diff --git a/docs/running/access_spectra.rst b/docs/running/access_spectra.rst
diff --git a/docs/running/index.rst b/docs/running/index.rst
@@ -0,0 +1,11 @@
+**************
+Running Tardis
+**************
+
+Information regarding running and operating Tardis.
+
+.. toctree::
+
+  ../configuration/index
+  access
+
diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py
@@ -13,7 +13,67 @@
 logger = logging.getLogger(__name__)
 
 
-class Simulation(HDFWriterMixin):
+class PlasmaStateStorerMixin(object):
+    """Mixin class to provide the capability to the simulation object of
+    storing plasma information and the inner boundary temperature during each
+    MC iteration.
+
+    Currently, storage for the dilution factor, the radiation temperature and
+    the electron density in each cell is provided. Additionally, the
+    temperature at the inner boundary is saved.
+    """
+    def __init__(self, iterations, no_of_shells):
+
+        self.iterations_w = np.zeros(
+            (iterations, no_of_shells))
+        self.iterations_t_rad = np.zeros(
+            (iterations, no_of_shells)) * u.K
+        self.iterations_electron_densities = np.zeros(
+            (iterations, no_of_shells))
+        self.iterations_t_inner = np.zeros(iterations) * u.K
+
+    def store_plasma_state(self, i, w, t_rad, electron_densities, t_inner):
+        """Store current plasma information and inner boundary temperature
+        used in iterated i.
+
+        Parameters
+        ----------
+        i : int
+            current iteration index (0 for the first)
+        w : np.ndarray
+            dilution factor
+        t_rad : astropy.units.Quantity
+            radiation temperature
+        electron_densities : np.ndarray
+            electron density
+        t_inner : astropy.units.Quantity
+            temperature of inner boundary
+        """
+        self.iterations_w[i, :] = w
+        self.iterations_t_rad[i, :] = t_rad
+        self.iterations_electron_densities[i, :] = \
+            electron_densities.values
+        self.iterations_t_inner[i] = t_inner
+
+    def reshape_plasma_state_store(self, executed_iterations):
+        """Reshapes the storage arrays in case convergence was reached before
+        all specified iterations were executed.
+
+        Parameters
+        ----------
+        executed_iterations : int
+            iteration index, i.e. number of iterations executed minus one!
+        """
+        self.iterations_w = self.iterations_w[:executed_iterations+1, :]
+        self.iterations_t_rad = \
+            self.iterations_t_rad[:executed_iterations+1, :]
+        self.iterations_electron_densities = \
+            self.iterations_electron_densities[:executed_iterations+1, :]
+        self.iterations_t_inner = \
+            self.iterations_t_inner[:executed_iterations+1]
+
+
+class Simulation(PlasmaStateStorerMixin, HDFWriterMixin):
     """A composite object containing all the required information for a
     simulation.
 
@@ -36,13 +96,18 @@ class Simulation(HDFWriterMixin):
         .. note:: TARDIS must be built with OpenMP support in order for
         `nthreads` to have effect.
     """
-    hdf_properties = ['model', 'plasma', 'runner']
+    hdf_properties = ['model', 'plasma', 'runner', 'iterations_w',
+                      'iterations_t_rad', 'iterations_electron_densities',
+                      'iterations_t_inner']
     hdf_name = 'simulation'
     def __init__(self, iterations, model, plasma, runner,
                  no_of_packets, no_of_virtual_packets, luminosity_nu_start,
                  luminosity_nu_end, last_no_of_packets,
                  luminosity_requested, convergence_strategy,
                  nthreads):
+
+        super(Simulation, self).__init__(iterations, model.no_of_shells)
+
         self.converged = False
         self.iterations = iterations
         self.iterations_executed = 0
@@ -217,20 +282,31 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0, last_run=False):
     def run(self):
         start_time = time.time()
         while self.iterations_executed < self.iterations-1:
+            self.store_plasma_state(self.iterations_executed, self.model.w,
+                                    self.model.t_rad,
+                                    self.plasma.electron_densities,
+                                    self.model.t_inner)
             self.iterate(self.no_of_packets)
             self.converged = self.advance_state()
             self._call_back()
             if self.converged:
                 if self.convergence_strategy.stop_if_converged:
                     break
         # Last iteration
+        self.store_plasma_state(self.iterations_executed, self.model.w,
+                                self.model.t_rad,
+                                self.plasma.electron_densities,
+                                self.model.t_inner)
         self.iterate(self.last_no_of_packets, self.no_of_virtual_packets, True)
 
+        self.reshape_plasma_state_store(self.iterations_executed)
+
         logger.info("Simulation finished in {0:d} iterations "
                     "and took {1:.2f} s".format(
                         self.iterations_executed, time.time() - start_time))
         self._call_back()
 
+
     def log_plasma_state(self, t_rad, w, t_inner, next_t_rad, next_w,
                          next_t_inner, log_sampling=5):
         """

diff --git a/tardis/simulation/tests/test_simulation.py b/tardis/simulation/tests/test_simulation.py
@@ -4,8 +4,10 @@
 from tardis.io.config_reader import Configuration
 from tardis.simulation import Simulation
 
+import numpy as np
 import pandas as pd
 import pandas.util.testing as pdt
+import astropy.units as u
 
 
 @pytest.fixture(scope='module')
@@ -37,6 +39,12 @@ def simulation_one_loop(
     if not generate_reference:
         return simulation
     else:
+        simulation.hdf_properties = [
+            'iterations_w',
+            'iterations_t_rad',
+            'iterations_electron_densities',
+            'iterations_t_inner',
+        ]
         simulation.model.hdf_properties = [
                 't_radiative',
                 'dilution_factor'
@@ -47,6 +55,11 @@ def simulation_one_loop(
                 'output_nu',
                 'output_energy'
                 ]
+        simulation.to_hdf(
+                tardis_ref_data,
+                '',
+                'test_simulation'
+        )
         simulation.model.to_hdf(
                 tardis_ref_data,
                 '',
@@ -79,5 +92,66 @@ def test_plasma_estimates(
             refdata(name)
             )
 
+
+@pytest.mark.parametrize('name', [
+    'iterations_w', 'iterations_t_rad',
+    'iterations_electron_densities', 'iterations_t_inner'
+    ])
+def test_plasma_state_iterations(
+        simulation_one_loop, refdata, name):
+    actual = getattr(
+        simulation_one_loop, name)
+
+    try:
+        actual = pd.Series(actual)
+    except Exception:
+        actual = pd.DataFrame(actual)
+
+    pdt.assert_almost_equal(
+            actual,
+            refdata(name)
+            )
+
+
+@pytest.fixture(scope="module")
+def simulation_without_loop(atomic_data_fname, config):
+
+    config.atom_data = atomic_data_fname
+    config.montecarlo.iterations = 2
+    return Simulation.from_config(config)
+
+
+def test_plasma_state_storer_store(atomic_data_fname, config,
+                                   simulation_without_loop):
+
+    simulation = simulation_without_loop
+
+    w_test = np.linspace(0, 1, 20)
+    t_rad_test = np.linspace(12000, 9000, 20) * u.K
+    electron_densities_test = pd.Series(np.linspace(1e7, 1e6, 20))
+    t_inner_test = 12500 * u.K
+
+    simulation.store_plasma_state(1, w_test, t_rad_test,
+                                  electron_densities_test, t_inner_test)
+
+    np.testing.assert_allclose(simulation.iterations_w[1, :], w_test)
+    np.testing.assert_allclose(simulation.iterations_t_rad[1, :], t_rad_test)
+    np.testing.assert_allclose(simulation.iterations_electron_densities[1, :],
+                               electron_densities_test)
+    np.testing.assert_allclose(simulation.iterations_t_inner[1], t_inner_test)
+
+
+def test_plasma_state_storer_reshape(atomic_data_fname, config,
+                                     simulation_without_loop):
+
+    simulation = simulation_without_loop
+    simulation.reshape_plasma_state_store(0)
+
+    assert simulation.iterations_t_rad.shape == (1, 20)
+    assert simulation.iterations_w.shape == (1, 20)
+    assert simulation.iterations_electron_densities.shape == (1, 20)
+    assert simulation.iterations_t_inner.shape == (1,)
+
+
 #     assert_quantity_allclose(
 #             t_rad, simulation_compare_data['test1/t_rad'] * u.Unit('K'), atol=0.0 * u.Unit('K'))