Merge ffd8b38 into da5664a

tardis/simulation/base.py

@@ -71,6 +71,15 @@ def __init__(self, iterations, model, plasma, runner,
                     'not damped or custom '
                     '- input is {0}'.format(convergence_strategy.type))
 
+        # containers to store plasma state (Tr, W, ne) for each iteration
+        self.iterations_w = np.zeros(
+            (self.iterations, self.model.no_of_shells))
+        self.iterations_t_rad = np.zeros(
+            (self.iterations, self.model.no_of_shells)) * u.K
+        self.iterations_t_inner = np.zeros(self.iterations) * u.K
+        self.iterations_electron_densities = np.zeros(
+            (self.iterations, self.model.no_of_shells))
+
         self._callbacks = OrderedDict()
         self._cb_next_id = 0
 
@@ -217,20 +226,40 @@ 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.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.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 store_plasma_state(self, i):
+
+        self.iterations_w[i, :] = self.model.w
+        self.iterations_t_rad[i, :] = self.model.t_rad
+        self.iterations_electron_densities[i, :] = \
+            self.plasma.electron_densities.values
+        self.iterations_t_inner[i] = self.model.t_inner
+
+    def reshape_plasma_state_store(self, executed_iterations):
+
+        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]
+
     def log_plasma_state(self, t_rad, w, t_inner, next_t_rad, next_w,
                          next_t_inner, log_sampling=5):
         """