Merge pull request #360 from tardis-sn/rebasedvirtrunner

Rebasedvirtrunner

tardis/model.py

@@ -9,9 +9,11 @@
 from astropy import constants, units as u
 import scipy.special
 
+from util import intensity_black_body
 from tardis import packet_source, plasma_array
 from tardis.montecarlo import montecarlo
-from util import intensity_black_body
+from tardis.montecarlo.base import MontecarloRunner
+
 
 
 logger = logging.getLogger(__name__)
@@ -20,12 +22,6 @@
 h = constants.h.cgs.value
 kb = constants.k_B.cgs.value
 
-w_estimator_constant = (c ** 2 / (2 * h)) * (15 / np.pi ** 4) * (h / kb) ** 4 / (4 * np.pi)
-
-t_rad_estimator_constant = (np.pi**4 / (15 * 24 * scipy.special.zeta(5, 1))) * h / kb
-
-
-
 
 class Radial1DModel(object):
     """
@@ -139,6 +135,7 @@ def __init__(self, tardis_config):
         self.spectrum = TARDISSpectrum(tardis_config.spectrum.frequency, tardis_config.supernova.distance)
         self.spectrum_virtual = TARDISSpectrum(tardis_config.spectrum.frequency, tardis_config.supernova.distance)
         self.spectrum_reabsorbed = TARDISSpectrum(tardis_config.spectrum.frequency, tardis_config.supernova.distance)
+        self.runner = MontecarloRunner()
 
 
 
@@ -202,38 +199,6 @@ def calculate_j_blues(self, init_detailed_j_blues=False):
         else:
             raise ValueError('radiative_rates_type type unknown - %s', radiative_rates_type)
 
-
-
-
-    def calculate_updated_radiationfield(self, nubar_estimator, j_estimator):
-        """
-        Calculate an updated radiation field from the :math:`\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}`
-        calculated in the montecarlo simulation. The details of the calculation can be found in the documentation.
-
-        Parameters
-        ----------
-
-        nubar_estimator : ~np.ndarray (float)
-
-        j_estimator : ~np.ndarray (float)
-
-        Returns
-        -------
-
-        updated_t_rads : ~np.ndarray (float)
-
-        updated_ws : ~np.ndarray (float)
-
-        """
-
-
-        updated_t_rads = t_rad_estimator_constant * nubar_estimator / j_estimator
-        updated_ws = j_estimator / (
-            4 * constants.sigma_sb.cgs.value * updated_t_rads ** 4 * self.time_of_simulation.value
-            * self.tardis_config.structure.volumes.value)
-
-        return updated_t_rads * u.K, updated_ws
-
     def update_plasmas(self, initialize_nlte=False):
 
         self.plasma_array.update_radiationfield(self.t_rads.value, self.ws, j_blues=self.j_blues,
@@ -249,7 +214,8 @@ def update_radiationfield(self, log_sampling=5):
         Updating radiation field
         """
         convergence_section = self.tardis_config.montecarlo.convergence_strategy
-        updated_t_rads, updated_ws = self.calculate_updated_radiationfield(self.nubar_estimators, self.j_estimators)
+        updated_t_rads, updated_ws = (
+            self.runner.calculate_radiationfield_properties())
         old_t_rads = self.t_rads.copy()
         old_ws = self.ws.copy()
         old_t_inner = self.t_inner
@@ -352,29 +318,35 @@ def simulate(self, update_radiation_field=True, enable_virtual=False, initialize
 
         self.j_blue_estimators = np.zeros((len(self.t_rads), len(self.atom_data.lines)))
         self.montecarlo_virtual_luminosity = np.zeros_like(self.spectrum.frequency.value)
-
-        montecarlo_nu, montecarlo_energies, self.j_estimators, self.nubar_estimators, \
-        last_line_interaction_in_id, last_line_interaction_out_id, \
-        self.last_interaction_type, self.last_line_interaction_shell_id = \
-            montecarlo.montecarlo_radial1d(self,
-                                                     virtual_packet_flag=no_of_virtual_packets, nthreads=self.tardis_config.montecarlo.nthreads)
+
+        self.runner.run(self, no_of_virtual_packets=no_of_virtual_packets,
+                        nthreads=self.tardis_config.montecarlo.nthreads) #self = model
+
+
+        (montecarlo_nu, montecarlo_energies, self.j_estimators,
+         self.nubar_estimators, last_line_interaction_in_id,
+         last_line_interaction_out_id, self.last_interaction_type,
+         self.last_line_interaction_shell_id) = self.runner.legacy_return()
 
         if np.sum(montecarlo_energies < 0) == len(montecarlo_energies):
             logger.critical("No r-packet escaped through the outer boundary.")
 
-        self.montecarlo_nu = montecarlo_nu * u.Hz
-        self.montecarlo_luminosity = montecarlo_energies *  1 * u.erg / self.time_of_simulation
+        self.montecarlo_nu = self.runner.packet_nu
+        self.montecarlo_luminosity = self.runner.packet_luminosity
+
+
+
+        montecarlo_reabsorbed_luminosity = np.histogram(
+            self.runner.reabsorbed_packet_nu,
+            weights=self.runner.reabsorbed_packet_luminosity,
+            bins=self.tardis_config.spectrum.frequency.value)[0] * u.erg / u.s
 
 
-        montecarlo_reabsorbed_luminosity = -np.histogram(self.montecarlo_nu.value[self.montecarlo_luminosity.value < 0],
-                                         weights=self.montecarlo_luminosity.value[self.montecarlo_luminosity.value < 0],
-                                         bins=self.tardis_config.spectrum.frequency.value)[0] \
-                                      * self.montecarlo_luminosity.unit
 
-        montecarlo_emitted_luminosity = np.histogram(self.montecarlo_nu.value[self.montecarlo_luminosity.value >= 0],
-                                         weights=self.montecarlo_luminosity.value[self.montecarlo_luminosity.value >= 0],
-                                         bins=self.tardis_config.spectrum.frequency.value)[0] \
-                                   * self.montecarlo_luminosity.unit
+        montecarlo_emitted_luminosity = np.histogram(
+            self.runner.emitted_packet_nu,
+            weights=self.runner.emitted_packet_luminosity,
+            bins=self.tardis_config.spectrum.frequency.value)[0] * u.erg / u.s
 
 
 

tardis/montecarlo/base.py

@@ -0,0 +1,104 @@
+from astropy import units as u, constants as const
+
+from scipy.special import zeta
+
+from tardis.montecarlo import montecarlo
+
+import numpy as np
+
+class MontecarloRunner(object):
+    """
+    This class is designed as an interface between the Python part and the
+    montecarlo C-part
+    """
+
+    w_estimator_constant = ((const.c ** 2 / (2 * const.h)) *
+                            (15 / np.pi ** 4) * (const.h / const.k_B) ** 4 /
+                            (4 * np.pi)).cgs.value
+
+    t_rad_estimator_constant = ((np.pi**4 / (15 * 24 * zeta(5, 1))) *
+                                (const.h / const.k_B)).cgs.value
+
+
+    def run(self, model, no_of_virtual_packets, nthreads=1):
+        self.time_of_simulation = model.time_of_simulation
+        self.volume = model.tardis_config.structure.volumes
+
+        montecarlo.montecarlo_radial1d(
+            model, self, virtual_packet_flag=no_of_virtual_packets,
+            nthreads=nthreads)
+
+    def legacy_return(self):
+        return (self.packet_nu, self.packet_energy,
+                self.j_estimator, self.nu_bar_estimator,
+                self.last_line_interaction_in_id,
+                self.last_line_interaction_out_id,
+                self.last_interaction_type,
+                self.last_line_interaction_shell_id)
+
+    @property
+    def packet_nu(self):
+        return u.Quantity(self._packet_nu, u.Hz)
+
+    @property
+    def packet_energy(self):
+        return u.Quantity(self._packet_energy, u.erg)
+
+    @property
+    def packet_luminosity(self):
+        return self.packet_energy / self.time_of_simulation
+
+    @property
+    def emitted_packet_mask(self):
+        return self.packet_energy >=0
+
+    @property
+    def emitted_packet_nu(self):
+        return self.packet_nu[self.emitted_packet_mask]
+
+    @property
+    def reabsorbed_packet_nu(self):
+        return self.packet_nu[~self.emitted_packet_mask]
+
+    @property
+    def reabsorbed_packet_luminosity(self):
+        return -self.packet_luminosity[~self.emitted_packet_mask]
+
+
+    @property
+    def emitted_packet_luminosity(self):
+        return self.packet_luminosity[self.emitted_packet_mask]
+
+    @property
+    def reabsorbed_packet_luminosity(self):
+        return -self.packet_luminosity[~self.emitted_packet_mask]
+
+    def calculate_radiationfield_properties(self):
+        """
+        Calculate an updated radiation field from the :math:`\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}`
+        calculated in the montecarlo simulation. The details of the calculation can be found in the documentation.
+
+        Parameters
+        ----------
+
+        nubar_estimator : ~np.ndarray (float)
+
+        j_estimator : ~np.ndarray (float)
+
+        Returns
+        -------
+
+        t_rad : ~astropy.units.Quantity (float)
+
+        w : ~numpy.ndarray (float)
+
+        """
+
+
+        t_rad = (self.t_rad_estimator_constant * self.nu_bar_estimator
+                / self.j_estimator)
+        w = self.j_estimator / (4 * const.sigma_sb.cgs.value * t_rad ** 4
+                                * self.time_of_simulation.value
+                                * self.volume.value)
+
+        return t_rad * u.K, w

tardis/montecarlo/montecarlo.pyx

@@ -11,6 +11,7 @@ import numpy as np
 cimport numpy as np
 from astropy import constants
 from astropy import units
+from libc.stdlib cimport free
 
 np.import_array()
 
@@ -73,10 +74,15 @@ cdef extern from "src/cmontecarlo.h":
         double *l_pop
         double *l_pop_r
         ContinuumProcessesStatus cont_status
+        double *virt_packet_nus
+        double *virt_packet_energies
+        int_type_t virt_packet_count
+        int_type_t virt_array_size
 
     void montecarlo_main_loop(storage_model_t * storage, int_type_t virtual_packet_flag, int nthreads, unsigned long seed)
 
-def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0, int nthreads=4):
+def montecarlo_radial1d(model, runner, int_type_t virtual_packet_flag=0,
+                        int nthreads=4):
     """
     Parameters
     ----------
@@ -222,5 +228,25 @@ def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0, int nthreads=4)
     #cdef np.ndarray[double, ndim=1] output_nus = np.zeros(storage.no_of_packets, dtype=np.float64)
     #cdef np.ndarray[double, ndim=1] output_energies = np.zeros(storage.no_of_packets, dtype=np.float64)
     montecarlo_main_loop(&storage, virtual_packet_flag, nthreads, model.tardis_config.montecarlo.seed)
-    return output_nus, output_energies, js, nubars, last_line_interaction_in_id, last_line_interaction_out_id, last_interaction_type, last_line_interaction_shell_id
+
+    cdef np.ndarray[double, ndim=1] virt_packet_nus = np.zeros(storage.virt_packet_count, dtype=np.float64)
+    cdef np.ndarray[double, ndim=1] virt_packet_energies = np.zeros(storage.virt_packet_count, dtype=np.float64)
+    for i in range(storage.virt_packet_count):
+        virt_packet_nus[i] = storage.virt_packet_nus[i]
+        virt_packet_energies[i] = storage.virt_packet_energies[i]
+    free(<void *>storage.virt_packet_nus)
+    free(<void *>storage.virt_packet_energies)
+    runner._packet_nu = output_nus
+    runner._packet_energy = output_energies
+    runner.j_estimator = js
+    runner.nu_bar_estimator = nubars
+    runner.last_line_interaction_in_id = last_line_interaction_in_id
+    runner.last_line_interaction_out_id = last_line_interaction_out_id
+    runner.last_interaction_type = last_interaction_type
+    runner.last_line_interaction_shell_id = last_line_interaction_shell_id
+    runner.virt_packet_nus = virt_packet_nus
+    runner.virt_packet_energies = virt_packet_energies
+
+    #return output_nus, output_energies, js, nubars, last_line_interaction_in_id, last_line_interaction_out_id, last_interaction_type, last_line_interaction_shell_id, virt_packet_nus, virt_packet_energies
+
 

tardis/montecarlo/src/cmontecarlo.c

@@ -479,12 +479,28 @@ montecarlo_one_packet (storage_model_t * storage, rpacket_t * packet,
 	      if ((virt_packet.nu < storage->spectrum_end_nu) &&
 		  (virt_packet.nu > storage->spectrum_start_nu))
 		{
-		  virt_id_nu =
-		    floor ((virt_packet.nu -
-			    storage->spectrum_start_nu) /
-			   storage->spectrum_delta_nu);
-		  storage->spectrum_virt_nu[virt_id_nu] +=
-		    virt_packet.energy * weight;
+#ifdef WITHOPENMP
+#pragma omp critical
+		  {
+#endif
+		    if (storage->virt_packet_count >= storage->virt_array_size)
+		      {
+			storage->virt_array_size *= 2;
+			storage->virt_packet_nus = realloc(storage->virt_packet_nus, sizeof(double) * storage->virt_array_size);
+			storage->virt_packet_energies = realloc(storage->virt_packet_energies, sizeof(double) * storage->virt_array_size);
+		      }
+		    storage->virt_packet_nus[storage->virt_packet_count] = virt_packet.nu;
+		    storage->virt_packet_energies[storage->virt_packet_count] = virt_packet.energy * weight;
+		    storage->virt_packet_count += 1;
+		    virt_id_nu =
+		      floor ((virt_packet.nu -
+			      storage->spectrum_start_nu) /
+			     storage->spectrum_delta_nu);
+		    storage->spectrum_virt_nu[virt_id_nu] +=
+		      virt_packet.energy * weight;
+#ifdef WITHOPENMP
+		  }
+#endif
 		}
 	    }
 	}
@@ -849,6 +865,10 @@ void
 montecarlo_main_loop(storage_model_t * storage, int64_t virtual_packet_flag, int nthreads, unsigned long seed)
 {
   int64_t packet_index;
+  storage->virt_packet_nus = (double *)malloc(sizeof(double) * storage->no_of_packets);
+  storage->virt_packet_energies = (double *)malloc(sizeof(double) * storage->no_of_packets);
+  storage->virt_packet_count = 0;
+  storage->virt_array_size = storage->no_of_packets;
 #ifdef WITHOPENMP
   fprintf(stderr, "Running with OpenMP - %d threads", nthreads);
   omp_set_dynamic(0);

tardis/montecarlo/src/storage.h

@@ -68,6 +68,10 @@ typedef struct StorageModel
   double *l_pop;
   double *l_pop_r;
   ContinuumProcessesStatus cont_status;
+  double *virt_packet_nus;
+  double *virt_packet_energies;
+  int64_t virt_packet_count;
+  int64_t virt_array_size;
 } storage_model_t;
 
 #endif // TARDIS_STORAGE_H