Merge pull request #595 from Tobychev/Lucy345

Line absorption rate estimator

tardis/model.py

@@ -123,6 +123,7 @@ def __init__(self, tardis_config):
             v_outer=self.v_outer)
 
         self.calculate_j_blues(init_detailed_j_blues=True)
+        self.Edotlu = np.zeros(np.shape(self.j_blues.shape))
         self.update_plasmas(initialize_nlte=True)
 
     @property
@@ -186,6 +187,7 @@ def t_inner(self, value):
              self.tardis_config.structure.volumes))
 
 
+
     @staticmethod
     def calculate_geometric_w(r, r_inner):
         return 0.5 * (1 - np.sqrt(1 - (r_inner ** 2 / r ** 2).to(1).value))

tardis/montecarlo/base.py

@@ -50,6 +50,7 @@ def _initialize_estimator_arrays(self, no_of_shells, tau_sobolev_shape):
         self.j_estimator = np.zeros(no_of_shells, dtype=np.float64)
         self.nu_bar_estimator = np.zeros(no_of_shells, dtype=np.float64)
         self.j_blue_estimator = np.zeros(tau_sobolev_shape)
+        self.Edotlu_estimator = np.zeros(tau_sobolev_shape)
 
 
     def _initialize_geometry_arrays(self, structure):
@@ -109,7 +110,7 @@ def legacy_update_spectrum(self, no_of_virtual_packets):
             self.spectrum_virtual.update_luminosity(
                 self.montecarlo_virtual_luminosity)
 
-    def run(self, model, no_of_packets, no_of_virtual_packets=0, nthreads=1):
+    def run(self, model, no_of_packets, no_of_virtual_packets=0, nthreads=1,last_run=False):
         """
         Running the TARDIS simulation
 
@@ -132,12 +133,14 @@ def run(self, model, no_of_packets, no_of_virtual_packets=0, nthreads=1):
 
         montecarlo.montecarlo_radial1d(
             model, self, virtual_packet_flag=no_of_virtual_packets,
-            nthreads=nthreads)
+            nthreads=nthreads,last_run=last_run)
         # Workaround so that j_blue_estimator is in the right ordering
         # They are written as an array of dimension (no_of_shells, no_of_lines)
         # but python expects (no_of_lines, no_of_shells)
         self.j_blue_estimator = self.j_blue_estimator.flatten().reshape(
                 self.j_blue_estimator.shape, order='F')
+        self.Edotlu_estimator = self.Edotlu_estimator.flatten().reshape(
+                self.Edotlu_estimator.shape, order='F')
 
     def legacy_return(self):
         return (self.output_nu, self.output_energy,

tardis/montecarlo/montecarlo.pyx

@@ -50,6 +50,7 @@ cdef extern from "src/cmontecarlo.h":
         double *continuum_list_nu
         int_type_t line_lists_tau_sobolevs_nd
         double *line_lists_j_blues
+        double *line_lists_Edotlu
         int_type_t line_lists_j_blues_nd
         int_type_t no_of_lines
         int_type_t no_of_edges
@@ -156,6 +157,9 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     storage.line_lists_j_blues = <double*> PyArray_DATA(
             runner.j_blue_estimator)
 
+    storage.line_lists_Edotlu = <double*> PyArray_DATA(
+            runner.Edotlu_estimator)
+
     storage.line_interaction_id = runner.get_line_interaction_id(
         model.tardis_config.plasma.line_interaction_type)
 
@@ -218,7 +222,7 @@ cdef initialize_storage_model(model, runner, storage_model_t *storage):
     storage.t_electrons = <double*> t_electrons.data
 
 def montecarlo_radial1d(model, runner, int_type_t virtual_packet_flag=0,
-                        int nthreads=4):
+                        int nthreads=4,last_run=False):
     """
     Parameters
     ----------
@@ -288,3 +292,14 @@ def montecarlo_radial1d(model, runner, int_type_t virtual_packet_flag=0,
         runner.virt_packet_last_interaction_type = np.zeros(0)
         runner.virt_packet_last_line_interaction_in_id = np.zeros(0)
         runner.virt_packet_last_line_interaction_out_id = np.zeros(0)
+
+    if last_run:
+        postprocess(model,runner)
+
+def postprocess(model, runner):
+    Edotlu_norm_factor = (1 / 
+        (model.time_of_simulation * model.tardis_config.structure.volumes))
+    exptau = 1 - np.exp(- 
+                        runner.line_lists_tau_sobolevs.reshape(-1,
+                            model.tardis_config.structure["velocity"]["num"]) ) 
+    model.Edotlu = Edotlu_norm_factor*exptau*runner.Edotlu_estimator        

tardis/montecarlo/src/cmontecarlo.c

@@ -373,6 +373,19 @@ increment_j_blue_estimator (const rpacket_t * packet, storage_model_t * storage,
     }
 }
 
+void
+increment_Edotlu_estimator (const rpacket_t * packet, storage_model_t * storage, int64_t line_idx)
+{
+  if (storage->line_lists_Edotlu != NULL)
+    {
+#ifdef WITHOPENMP
+#pragma omp atomic
+#endif
+    storage->line_lists_Edotlu[line_idx] += rpacket_get_energy (packet);
+    }
+}
+
+
 int64_t
 montecarlo_one_packet (storage_model_t * storage, rpacket_t * packet,
                        int64_t virtual_mode, rk_state *mt_state)
@@ -613,6 +626,7 @@ montecarlo_line_scatter (rpacket_t * packet, storage_model_t * storage,
   if (rpacket_get_virtual_packet (packet) == 0)
     {
       increment_j_blue_estimator (packet, storage, distance, line2d_idx);
+      increment_Edotlu_estimator (packet, storage, line2d_idx);
     }
   double tau_line =
     storage->line_lists_tau_sobolevs[line2d_idx];

tardis/montecarlo/src/cmontecarlo.h

@@ -67,6 +67,10 @@ void increment_j_blue_estimator (const rpacket_t * packet,
                                  storage_model_t * storage,
                                  double d_line, int64_t j_blue_idx);
 
+void increment_Edotlu_estimator (const rpacket_t * packet,
+                                 storage_model_t * storage,
+                                 int64_t j_blue_idx);
+
 int64_t montecarlo_one_packet (storage_model_t * storage, rpacket_t * packet,
                                int64_t virtual_mode, rk_state *mt_state);
 

tardis/montecarlo/src/storage.h

@@ -32,6 +32,7 @@ typedef struct StorageModel
   int64_t line_lists_tau_sobolevs_nd;
   double *line_lists_j_blues;
   int64_t line_lists_j_blues_nd;
+  double *line_lists_Edotlu;
   int64_t no_of_lines;
   int64_t no_of_edges;
   int64_t line_interaction_id;

tardis/montecarlo/struct.py

@@ -61,6 +61,7 @@ class StorageModel(Structure):
         ('line_lists_tau_sobolevs_nd', c_int64),
         ('line_lists_j_blues', POINTER(c_double)),
         ('line_lists_j_blues_nd', c_int64),
+        ('line_lists_Edotlu', POINTER(c_double)),
         ('no_of_lines', c_int64),
         ('no_of_edges', c_int64),
         ('line_interaction_id', c_int64),

tardis/montecarlo/tests/test_cmontecarlo.py

@@ -119,6 +119,8 @@ def model():
         line_lists_j_blues=(c_double * 2)(*([1.e-10] * 2)),
         line_lists_j_blues_nd=0,
 
+        line_lists_Edotlu=(c_double * 3)(*[0.0,0.0,1.0]), # Init to an explicit array 
+
         no_of_lines=2,
         no_of_edges=100,
 
@@ -426,7 +428,24 @@ def test_montecarlo_line_scatter(packet_params, expected_params, packet, model,
     assert_almost_equal(packet.tau_event, expected_params['tau_event'])
     assert_almost_equal(packet.next_line_id, expected_params['next_line_id'])
 
+"""
+Simple Tests:
+----------------
+These test check very simple pices of code still work.
+"""
+
+@pytest.mark.parametrize(
+    ['packet_params', 'line_idx', 'expected'],
+    [({'energy':0.0}, 0, 0),
+     ({'energy':1.0}, 1, 1),
+     ({'energy':0.5}, 2, 1.5)]
+)
+def test_increment_Edotlu_estimator(packet_params, line_idx, expected, packet, model):
+    packet.energy = packet_params['energy']
+
+    cmontecarlo_methods.increment_Edotlu_estimator(byref(packet), byref(model), c_int64(line_idx))
 
+    assert_almost_equal(model.line_lists_Edotlu[line_idx], expected)
 
 """
 Difficult Tests:

tardis/simulation/base.py

@@ -33,7 +33,7 @@ def __init__(self, tardis_config):
         self.t_inner_update = itertools.cycle(t_inner_lock_cycle)
 
     def run_single_montecarlo(self, model, no_of_packets,
-                              no_of_virtual_packets=0):
+                              no_of_virtual_packets=0,last_run=False):
         """
         Will do a single TARDIS iteration with the given model
         Parameters
@@ -51,7 +51,7 @@ def run_single_montecarlo(self, model, no_of_packets,
         """
         self.runner.run(model, no_of_packets,
                         no_of_virtual_packets=no_of_virtual_packets,
-                        nthreads=self.tardis_config.montecarlo.nthreads)
+                        nthreads=self.tardis_config.montecarlo.nthreads,last_run=last_run)
 
 
         (montecarlo_nu, montecarlo_energies, self.j_estimators,
@@ -322,10 +322,11 @@ def legacy_run_simulation(self, model, hdf_path_or_buf=None,
         no_of_virtual_packets = (
             self.tardis_config.montecarlo.no_of_virtual_packets)
 
-        self.run_single_montecarlo(model, no_of_packets, no_of_virtual_packets)
+        self.run_single_montecarlo(model, no_of_packets, no_of_virtual_packets, last_run=True)
 
         self.runner.legacy_update_spectrum(no_of_virtual_packets)
         self.legacy_set_final_model_properties(model)
+        model.Edotlu_estimators = self.runner.Edotlu_estimator
 
         #the following instructions, passing down information to the model are
         #required for the gui